Choose a GPU in main.py

This commit is contained in:
Steven Robertson 2017-04-21 13:08:16 -07:00
parent b507c9d604
commit c79db04490

74
main.py
View File

@ -118,42 +118,48 @@ def main(args, prof):
frames = [('%s%05d%s' % (prefix_plus, i, args.suffix), t) frames = [('%s%05d%s' % (prefix_plus, i, args.suffix), t)
for i, t in profile.enumerate_times(gprof)] for i, t in profile.enumerate_times(gprof)]
# We don't initialize a CUDA context until here. This keeps other import pycuda.driver as cuda
# functions like --help and --print snappy. cuda.init()
import pycuda.autoinit dev = cuda.Device(args.device or 0)
rmgr = render.RenderManager() cuctx = dev.make_context()
rdr = render.Renderer(gnm, gprof)
def render_iter(): try:
m = os.path.getmtime(args.flame) rmgr = render.RenderManager()
first = True rdr = render.Renderer(gnm, gprof)
for name, times in frames:
if args.resume:
fp = name + rdr.out.suffix
if os.path.isfile(fp) and m < os.path.getmtime(fp):
continue
for idx, t in enumerate(times): def render_iter():
evt, buf = rmgr.queue_frame(rdr, gnm, gprof, t, first) m = os.path.getmtime(args.flame)
first = False first = True
while not evt.query(): for name, times in frames:
time.sleep(0.01) if args.resume:
yield None fp = name + rdr.out.suffix
save(rdr.out, name, buf) if os.path.isfile(fp) and m < os.path.getmtime(fp):
if args.rawfn: continue
try:
buf.tofile(args.rawfn + '.tmp') for idx, t in enumerate(times):
os.rename(args.rawfn + '.tmp', args.rawfn) evt, buf = rmgr.queue_frame(rdr, gnm, gprof, t, first)
except e: first = False
print 'Failed to write %s: %s' % (args.rawfn, e) while not evt.query():
print '%s (%3d/%3d), %dms' % (name, idx, len(times), evt.time()) time.sleep(0.01)
yield name, buf yield None
save(rdr.out, name, None) save(rdr.out, name, buf)
if args.rawfn:
try:
buf.tofile(args.rawfn + '.tmp')
os.rename(args.rawfn + '.tmp', args.rawfn)
except e:
print 'Failed to write %s: %s' % (args.rawfn, e)
print '%s (%3d/%3d), %dms' % (name, idx, len(times), evt.time())
yield name, buf
save(rdr.out, name, None)
if args.gfx:
pyglet_preview(args, gprof, render_iter())
else:
for i in render_iter(): pass
finally:
cuda.Context.pop()
if args.gfx:
pyglet_preview(args, gprof, render_iter())
else:
for i in render_iter(): pass
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Render fractal flames.') parser = argparse.ArgumentParser(description='Render fractal flames.')
@ -183,6 +189,8 @@ if __name__ == "__main__":
help='Use half-loops when converting nodes to animations') help='Use half-loops when converting nodes to animations')
parser.add_argument('--print', action='store_true', parser.add_argument('--print', action='store_true',
help="Print the blended animation and exit.") help="Print the blended animation and exit.")
parser.add_argument('--device', metavar='NUM', type=int,
help="GPU device number to use (from nvidia-smi).")
profile.add_args(parser) profile.add_args(parser)
args = parser.parse_args() args = parser.parse_args()