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

14
main.py
View File

@ -118,9 +118,12 @@ 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)
cuctx = dev.make_context()
try:
rmgr = render.RenderManager() rmgr = render.RenderManager()
rdr = render.Renderer(gnm, gprof) rdr = render.Renderer(gnm, gprof)
@ -154,6 +157,9 @@ def main(args, prof):
pyglet_preview(args, gprof, render_iter()) pyglet_preview(args, gprof, render_iter())
else: else:
for i in render_iter(): pass for i in render_iter(): pass
finally:
cuda.Context.pop()
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()