Update app.py

app.py

@@ -52,7 +52,7 @@ current_model = models[1] if is_colab else models[0]
 current_model_path = current_model.path
 
 if is_colab:
-  pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
+  pipe = StableDiffusionPipeline.from_pretrained(current_model.path, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False)) #add torch_dtype=torch.float16 for GPU
 
 else: # download all models
   print(f"{datetime.datetime.now()} Downloading vae...")
@@ -61,8 +61,8 @@ else: # download all models
     try:
         print(f"{datetime.datetime.now()} Downloading {model.name} model...")
         unet = UNet2DConditionModel.from_pretrained(model.path, subfolder="unet", torch_dtype=torch.float16)
-        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
-        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, torch_dtype=torch.float16, scheduler=scheduler)
+        model.pipe_t2i = StableDiffusionPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler)
+        model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(model.path, unet=unet, vae=vae, scheduler=scheduler) #I could add another If/Else if you want so it works with both CPU and GPU... 
     except Exception as e:
         print(f"{datetime.datetime.now()} Failed to load model " + model.name + ": " + str(e))
         models.remove(model)
@@ -152,7 +152,7 @@ def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, w
         current_model_path = model_path
 
         if is_colab or current_model == custom_model:
-          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False))
+          pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, scheduler=scheduler, safety_checker=lambda images, clip_input: (images, False)) #removed another torch_dtype=torch.float16
         else:
           pipe = pipe.to("cpu")
           pipe = current_model.pipe_i2i