hongfz16 commited on
Commit
45d6b3c
1 Parent(s): 699ffc9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -60,7 +60,7 @@ elif ckpt.endswith(".safetensors"):
60
  model.load_state_dict(model_ckpt)
61
  else:
62
  raise NotImplementedError
63
- device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
64
  model = model.to(device)
65
  sampler = DDIMSampler(model)
66
 
@@ -116,7 +116,7 @@ opt.prompt = ''
116
  opt.text_dir = True
117
  opt.front_dir = '+z'
118
  opt.force_cuda_rast = True
119
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
120
  gui = GUI(opt)
121
  ###################################### INIT STAGE 2 #########################################
122
 
@@ -348,7 +348,7 @@ def process_stage2(input_model, input_text, input_dir, iters, output_model, outp
348
  opt.mesh = input_model
349
 
350
  # load mesh!
351
- gui.renderer = gui.renderer_class(opt, device).to(device)
352
 
353
  # set prompt
354
  gui.prompt = opt.positive_prompt + ', ' + input_text
 
60
  model.load_state_dict(model_ckpt)
61
  else:
62
  raise NotImplementedError
63
+ device = torch.device("cuda:1") if torch.cuda.is_available() else torch.device("cpu")
64
  model = model.to(device)
65
  sampler = DDIMSampler(model)
66
 
 
116
  opt.text_dir = True
117
  opt.front_dir = '+z'
118
  opt.force_cuda_rast = True
119
+ device0 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
120
  gui = GUI(opt)
121
  ###################################### INIT STAGE 2 #########################################
122
 
 
348
  opt.mesh = input_model
349
 
350
  # load mesh!
351
+ gui.renderer = gui.renderer_class(opt, device0).to(device0)
352
 
353
  # set prompt
354
  gui.prompt = opt.positive_prompt + ', ' + input_text