Update app.py

app.py

@@ -60,7 +60,7 @@ elif ckpt.endswith(".safetensors"):
     model.load_state_dict(model_ckpt)
 else:
     raise NotImplementedError
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+device = torch.device("cuda:1") if torch.cuda.is_available() else torch.device("cpu")
 model = model.to(device)
 sampler = DDIMSampler(model)
 
@@ -116,7 +116,7 @@ opt.prompt = ''
 opt.text_dir = True
 opt.front_dir = '+z'
 opt.force_cuda_rast = True
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device0 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 gui = GUI(opt)
 ###################################### INIT STAGE 2 #########################################
 
@@ -348,7 +348,7 @@ def process_stage2(input_model, input_text, input_dir, iters, output_model, outp
     opt.mesh = input_model
 
     # load mesh!
-    gui.renderer = gui.renderer_class(opt, device).to(device)
+    gui.renderer = gui.renderer_class(opt, device0).to(device0)
 
     # set prompt
     gui.prompt = opt.positive_prompt + ', ' + input_text