Update app.py

app.py

@@ -44,6 +44,7 @@ USE_TORCH_COMPILE = 0
 ENABLE_CPU_OFFLOAD = 0
 
 
+
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "fluently/Fluently-XL-v4",
@@ -58,41 +59,42 @@ if torch.cuda.is_available():
 
     pipe.to("cuda")
     
-
-def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-):
+    def generate(
+        prompt: str,
+        negative_prompt: str = "",
+        use_negative_prompt: bool = False,
+        seed: int = 0,
+        width: int = 1024,
+        height: int = 1024,
+        guidance_scale: float = 3,
+        randomize_seed: bool = False,
+    ):
+        seed = int(randomize_seed_fn(seed, randomize_seed))
+
+        if not use_negative_prompt:
+            negative_prompt = ""  # type: ignore
+
+        images = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=20,
+            num_images_per_prompt=1,
+            cross_attention_kwargs={"scale": 0.65},
+            output_type="pil",
+        ).images
+        image_paths = [save_image(img, prompt) for img in images]
+
+        download_links = [create_download_link(path) for path in image_paths]
+
+        print(image_paths)
+        return image_paths, seed, download_links
+else:
+    st.warning("CUDA is not available. The demo may not work on CPU.")
 
     
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-
-    if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
-
-    images = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=20,
-        num_images_per_prompt=1,
-        cross_attention_kwargs={"scale": 0.65},
-        output_type="pil",
-    ).images
-    image_paths = [save_image(img, prompt) for img in images]
-
-    download_links = [create_download_link(path) for path in image_paths]
-
-    print(image_paths)
-    return image_paths, seed, download_links
 
 examples = [
 "a modern hospital room with advanced medical equipment and a patient resting comfortably",