Update app.py

app.py

@@ -2,30 +2,29 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load model and specify device (GPU if available, otherwise CPU)
+# Load the model and specify device (GPU if available)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the pipeline with accelerate support for memory optimization
+# Load the pipeline
 pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2")
 pipe = pipe.to(device)
 
 def generate_comic_images(story_text):
-    # Split the story into sentences
-    scenes = story_text.split(".")
+    # Split the story into sentences (scenes)
+    scenes = [scene.strip() for scene in story_text.split(".") if scene.strip()]
     
-    # Generate an image for each scene using the pipeline
     comic_images = []
     for scene in scenes:
-        scene = scene.strip()
         if scene:
             print(f"Generating image for scene: {scene}")
             try:
-                # Lower number of steps for faster generation
+                # Generate the image for each scene
                 image = pipe(scene, num_inference_steps=20, height=512, width=512).images[0]
+                # Convert the image to a format Gradio can handle (like PIL or numpy array)
                 comic_images.append(image)
             except Exception as e:
                 print(f"Error generating image for '{scene}': {e}")
-                comic_images.append(f"Error generating image for '{scene}'")
+                comic_images.append(None)
     
     return comic_images
 
@@ -42,9 +41,10 @@ with gr.Blocks() as demo:
     
     generate_button = gr.Button("Generate Comic")
     
+    # Gallery component to display generated images
     comic_output = gr.Gallery(label="Generated Comic", show_label=True, elem_id="comic-gallery")
     
     generate_button.click(fn=comic_generation_interface, inputs=story_input, outputs=comic_output)
 
-# Launch the app
+# Launch the Gradio app
 demo.launch()