Update app.py

app.py

@@ -2,47 +2,46 @@ import torch
 from diffusers import StableDiffusionPipeline
 import imageio
 import os
+import gradio as gr
 
-def get_prompt():
-    """Ensures user provides a valid prompt, avoiding EOFError."""
-    while True:
-        try:
-            prompt = input("Enter your prompt: ").strip()
-            if prompt:
-                return prompt
-            else:
-                print("Prompt cannot be empty. Please enter a valid prompt.")
-        except EOFError:
-            print("\nError: Input required. Restarting prompt entry...")
-            continue
+# Ensure the output directory exists
+OUTPUT_FOLDER = "generated_frames"
+os.makedirs(OUTPUT_FOLDER, exist_ok=True)
 
-def generate_images(prompt, num_images=5, output_folder="generated_frames"):
-    """Generates images using Stable Diffusion and creates an animated GIF."""
-    
-    # Ensure output folder exists
-    os.makedirs(output_folder, exist_ok=True)
-    
-    # Load model (CPU only)
-    print("Loading Stable Diffusion model...")
-    pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-    pipe.to("cpu")  # Force to run on CPU
+# Load Stable Diffusion Model
+print("Loading Stable Diffusion model...")
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+pipe.to("cpu")  # Force to run on CPU
 
+def generate_images(prompt, num_images=5):
+    """Generates images from a prompt and creates an animated GIF."""
+    
     images = []
     for i in range(num_images):
         print(f"Generating image {i+1}/{num_images}...")
         image = pipe(prompt).images[0]  # Generate image
-        image_path = os.path.join(output_folder, f"frame_{i}.png")
+        image_path = os.path.join(OUTPUT_FOLDER, f"frame_{i}.png")
         image.save(image_path)
-        images.append(imageio.v3.imread(image_path))  # Use v3 API to avoid deprecation warnings
+        images.append(imageio.v3.imread(image_path))  # Use v3 API
 
     # Create GIF animation
-    gif_path = os.path.join(output_folder, "animation.gif")
+    gif_path = os.path.join(OUTPUT_FOLDER, "animation.gif")
     imageio.mimsave(gif_path, images, duration=0.5)
     
-    print(f"\n✅ Animation saved at: {gif_path}")
-    return gif_path
+    return gif_path  # Return GIF path for download
 
-if __name__ == "__main__":
-    prompt = get_prompt()  # Get user input safely
-    generate_images(prompt)
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# Stable Diffusion Animation Generator 🎨✨")
+    
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt")
+        generate_button = gr.Button("Generate Animation 🎥")
+    
+    gif_output = gr.File(label="Download your GIF")
 
+    generate_button.click(fn=generate_images, inputs=prompt_input, outputs=gif_output)
+
+# Launch Gradio App
+if __name__ == "__main__":
+    demo.launch()