Update app.py

app.py

@@ -138,9 +138,13 @@ def txt_to_img(
 ):
     pipe = MODELS[model_name].pipe_t2i
 
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
-        pipe.enable_xformers_memory_efficient_attention()
+    if pipe is not None:
+        if torch.cuda.is_available():
+            pipe = pipe.to("cuda")
+            pipe.enable_xformers_memory_efficient_attention()
+			
+	else:
+        raise ValueError(f"Unable to find pipeline for model: {model_name}")
 
     result = pipe(
         prompt,
@@ -153,8 +157,9 @@ def txt_to_img(
         generator=generator,
     )
 
-    pipe.to("cpu")
-    torch.cuda.empty_cache()
+    if pipe is not None:
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
 
     return replace_nsfw_images(result)
 
@@ -175,9 +180,13 @@ def img_to_img(
 ):
     pipe = model.pipe_i2i
 
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
-        pipe.enable_xformers_memory_efficient_attention()
+    if pipe is not None:
+        if torch.cuda.is_available():
+            pipe = pipe.to("cuda")
+            pipe.enable_xformers_memory_efficient_attention()
+
+    else:
+        raise ValueError(f"Unable to find pipeline for model: {model_name}")
 
     ratio = min(height / img.height, width / img.width)
     img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
@@ -193,8 +202,9 @@ def img_to_img(
         generator=generator,
     )
 
-    pipe.to("cpu")
-    torch.cuda.empty_cache()
+    if pipe is not None:
+        pipe.to("cpu")
+        torch.cuda.empty_cache()
 
     return replace_nsfw_images(result)