Update app.py

app.py

@@ -1,32 +1,34 @@
 import gradio as gr
-import torch
 from diffusers import StableDiffusionImg2ImgPipeline
+import torch
 from PIL import Image
-import numpy as np
 
-# Load pre-trained model (no fine-tuning required)
+# Load pretrained stable-diffusion model
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
-).to("cuda")
+    "runwayml/stable-diffusion-v1-5", 
+    torch_dtype=torch.float16, 
+    revision="fp16", 
+    use_auth_token=True
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
 def generate_thumbnail(prompt, input_image):
-    if input_image is None or prompt.strip() == "":
-        return None
-
-    # Convert image to 512x512
-    input_image = input_image.resize((512, 512))
+    # Resize image to 512x512
+    image = input_image.resize((512, 512)).convert("RGB")
     
-    result = pipe(prompt=prompt, image=input_image, strength=0.75, guidance_scale=7.5).images[0]
-    return result
+    # Generate image
+    output = pipe(prompt=prompt, image=image, strength=0.75, guidance_scale=7.5)
+    return output.images[0]
 
-gr.Interface(
+# Gradio UI
+demo = gr.Interface(
     fn=generate_thumbnail,
     inputs=[
-        gr.Textbox(label="Prompt (e.g. 'NOOB vs PRO Minecraft epic battle')"),
-        gr.Image(label="Background Image", type="pil")
+        gr.Textbox(label="Enter your prompt (e.g., NOOB vs PRO Minecraft style)"),
+        gr.Image(type="pil", label="Upload Background Image"),
     ],
     outputs=gr.Image(label="Generated Thumbnail"),
-    title="🖼️ AI Thumbnail Generator (Stable Diffusion)",
-    description="Upload a background + write your prompt to generate a Minecraft-style thumbnail using Stable Diffusion v1.5"
-).launch()
+    title="AI Thumbnail Generator",
+    description="Upload an image and enter your text to create a Minecraft-style YouTube thumbnail"
+)
+
+demo.launch()