use cpu

app.py

@@ -58,10 +58,9 @@ current_model_path = current_model.path
 
 auth_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 
-if is_colab:
-  pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler, use_auth_token=auth_token)
+print(f"Is CUDA available: {torch.cuda.is_available()}")
 
-else: # download all models
+if torch.cuda.is_available():
   vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", torch_dtype=torch.float16, use_auth_token=auth_token)
   for model in models:
     try:
@@ -71,10 +70,20 @@ else: # download all models
     except:
         models.remove(model)
   pipe = models[0].pipe_t2i
-  
-if torch.cuda.is_available():
   pipe = pipe.to("cuda")
 
+else:
+  vae = AutoencoderKL.from_pretrained(current_model.path, subfolder="vae", use_auth_token=auth_token)
+  for model in models:
+    try:
+        unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", use_auth_token=auth_token)
+        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler, use_auth_token=auth_token)
+        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler, use_auth_token=auth_token)
+    except:
+        models.remove(model)
+  pipe = models[0].pipe_t2i
+  pipe = pipe.to("cpu")
+
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
 def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
@@ -100,11 +109,8 @@ def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, g
     if model_path != current_model_path or last_mode != "txt2img":
         current_model_path = model_path
 
-        if is_colab or current_model == models[0]:
-          pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, use_auth_token=auth_token)
-        else:
-          pipe.to("cpu")
-          pipe = current_model.pipe_t2i
+        pipe.to("cpu")
+        pipe = current_model.pipe_t2i
 
         if torch.cuda.is_available():
           pipe = pipe.to("cuda")
@@ -131,11 +137,8 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
     if model_path != current_model_path or last_mode != "img2img":
         current_model_path = model_path
 
-        if is_colab or current_model == models[0]:
-          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, use_auth_token=auth_token)
-        else:
-          pipe.to("cpu")
-          pipe = current_model.pipe_i2i
+        pipe.to("cpu")
+        pipe = current_model.pipe_i2i
         
         if torch.cuda.is_available():
           pipe = pipe.to("cuda")