Update app.py

app.py

@@ -58,22 +58,27 @@ models = [
   ]
 
 
-last_mode = "img2img"
+custom_model = None
+if is_colab:
+  models.insert(0, Model("Custom model"))
+  custom_model = models[0]
+
+last_mode = "txt2img"
 current_model = models[1] if is_colab else models[0]
 current_model_path = current_model.path
 
 if is_colab:
-   pipe = StableDiffusionPipeline.from_pretrained(
+  pipe = StableDiffusionPipeline.from_pretrained(
       current_model.path,
-      torch_dtype=torch.float16,
-      scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
+      torch_dtype=torch.float32,
+      scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
+      safety_checker=lambda images, clip_input: (images, False)
       )
 
-
 else:
-    pipe = StableDiffusionPipeline.from_pretrained(
+  pipe = StableDiffusionPipeline.from_pretrained(
       current_model.path,
-      torch_dtype=torch.float16,
+      torch_dtype=torch.float32,
       scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
       )
     
@@ -129,7 +134,7 @@ def inference(model_name, prompt, guidance, steps, n_images=1, width=512, height
   if seed == 0:
     seed = random.randint(0, 2147483647)
 
-  generator = torch.Generator('cuda').manual_seed(seed)
+  generator = torch.Generator('cpu').manual_seed(seed)
 
   try:
     if img is not None:
@@ -154,14 +159,14 @@ def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width,
         if is_colab or current_model == custom_model:
           pipe = StableDiffusionPipeline.from_pretrained(
               current_model_path,
-              torch_dtype=torch.float16,
+              torch_dtype=torch.float32,
               scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
               safety_checker=lambda images, clip_input: (images, False)
               )
         else:
           pipe = StableDiffusionPipeline.from_pretrained(
               current_model_path,
-              torch_dtype=torch.float16,
+              torch_dtype=torch.float32,
               scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
               )
           # pipe = pipe.to("cpu")
@@ -185,7 +190,7 @@ def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width,
       callback=pipe_callback)
 
     # update_state(f"Done. Seed: {seed}")
-
+    
     return replace_nsfw_images(result)
 
 def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed):
@@ -203,14 +208,14 @@ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance
         if is_colab or current_model == custom_model:
           pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
               current_model_path,
-              torch_dtype=torch.float16,
+              torch_dtype=torch.float32,
               scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"),
               safety_checker=lambda images, clip_input: (images, False)
               )
         else:
           pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
               current_model_path,
-              torch_dtype=torch.float16,
+              torch_dtype=torch.float32,
               scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler")
               )
           # pipe = pipe.to("cpu")
@@ -238,6 +243,8 @@ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance
         callback=pipe_callback)
 
     # update_state(f"Done. Seed: {seed}")
+        
+    return replace_nsfw_images(result)
 
 def replace_nsfw_images(results):