Update app.py

app.py

@@ -8,15 +8,15 @@ import torch
 from typing import Optional
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id_default = "CompVis/stable-diffusion-v1-4"
+model_id_default = "CompVis/stable-diffusion-v1-4"
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
 else:
     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id_default, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+pipe_default = DiffusionPipeline.from_pretrained(model_id_default, torch_dtype=torch_dtype)
+pipe_default = pipe_default.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -36,19 +36,24 @@ def infer(
 ):
     generator = torch.Generator().manual_seed(seed)
 
+    params = {
+        'prompt': prompt,
+        'negative_prompt': negative_prompt,
+        'guidance_scale': guidance_scale,
+        'num_inference_steps': num_inference_steps,
+        'width': width,
+        'height': height,
+        'generator': generator,
+    }
+
     if model_id != model_repo_id_default:
         pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
         pipe = pipe.to(device)
+        image = pipe(**params).images[0]
+    else:
+        image = pipe_default(**params).images[0]
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    
 
     return image, pipe.name_or_path