r3gm commited on
Commit
20d8039
1 Parent(s): 706551f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -367,6 +367,7 @@ class GuiSD:
367
  vae_model=None,
368
  type_model_precision=torch.float16,
369
  retain_task_model_in_cache=False,
 
370
  )
371
 
372
  def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
@@ -384,7 +385,8 @@ class GuiSD:
384
  if incompatible_vae:
385
  vae_model = None
386
 
387
-
 
388
  self.model.load_pipe(
389
  model_name,
390
  task_name=task_stablepy[task],
@@ -555,6 +557,7 @@ class GuiSD:
555
 
556
  # First load
557
  model_precision = torch.float16
 
558
  if not self.model:
559
  from modelstream import Model_Diffusers2
560
 
 
367
  vae_model=None,
368
  type_model_precision=torch.float16,
369
  retain_task_model_in_cache=False,
370
+ device="cpu",
371
  )
372
 
373
  def load_new_model(self, model_name, vae_model, task, progress=gr.Progress(track_tqdm=True)):
 
385
  if incompatible_vae:
386
  vae_model = None
387
 
388
+ self.model.device = torch.device("cpu")
389
+
390
  self.model.load_pipe(
391
  model_name,
392
  task_name=task_stablepy[task],
 
557
 
558
  # First load
559
  model_precision = torch.float16
560
+ self.model.device = torch.device("cuda:0")
561
  if not self.model:
562
  from modelstream import Model_Diffusers2
563