Update app.py

app.py

@@ -3,32 +3,33 @@ from diffusers import StableDiffusionImg2ImgPipeline
 import torch
 from PIL import Image
 
-# Load pretrained stable-diffusion model
+# Load the pretrained model
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", 
-    torch_dtype=torch.float16, 
-    revision="fp16", 
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,   # Or float32 if on CPU
     use_auth_token=True
-).to("cuda" if torch.cuda.is_available() else "cpu")
+)
+
+# Set to GPU or CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
 
+# Thumbnail generation function
 def generate_thumbnail(prompt, input_image):
-    # Resize image to 512x512
     image = input_image.resize((512, 512)).convert("RGB")
-    
-    # Generate image
-    output = pipe(prompt=prompt, image=image, strength=0.75, guidance_scale=7.5)
-    return output.images[0]
+    result = pipe(prompt=prompt, image=image, strength=0.75, guidance_scale=7.5)
+    return result.images[0]
 
 # Gradio UI
-demo = gr.Interface(
+interface = gr.Interface(
     fn=generate_thumbnail,
     inputs=[
-        gr.Textbox(label="Enter your prompt (e.g., NOOB vs PRO Minecraft style)"),
-        gr.Image(type="pil", label="Upload Background Image"),
+        gr.Textbox(label="Thumbnail Prompt (e.g., 'NOOB vs PRO in Minecraft')"),
+        gr.Image(type="pil", label="Upload Background Image")
     ],
     outputs=gr.Image(label="Generated Thumbnail"),
-    title="AI Thumbnail Generator",
-    description="Upload an image and enter your text to create a Minecraft-style YouTube thumbnail"
+    title="🖼️ AI Thumbnail Generator",
+    description="Upload a background image and enter a thumbnail prompt to generate your own YouTube-style thumbnail using Stable Diffusion!"
 )
 
-demo.launch()
+interface.launch()