Update app.py

app.py

@@ -1,42 +1,51 @@
 # -*- coding: utf-8 -*-
 """
-Gradio Space: Text → Image using Stable Diffusion (Hugging Face Inference API)
+Gradio Space: Text → Image (Diffusers Pipeline)
 UI designed by Mehak Mazhar
 """
 
 import os
-import io
-import requests
-from PIL import Image
+import torch
+from diffusers import StableDiffusionPipeline
 import gradio as gr
 
-# --- Configuration ---
-MODEL = os.environ.get("SD_MODEL", "runwayml/stable-diffusion-v1-5")
-HF_API_URL = f"https://api-inference.huggingface.co/models/{MODEL}"
+# --- Available models ---
+MODEL_CHOICES = {
+    "Dreamlike Diffusion 1.0": "dreamlike-art/dreamlike-diffusion-1.0",
+    "Stable Diffusion XL Base": "stabilityai/stable-diffusion-xl-base-1.0"
+}
 
-# --- API call function ---
-def generate_image(prompt, token, width, height, guidance_scale, steps):
-    if not token:
-        return None, "❌ Please provide a Hugging Face API token."
+# --- Cache pipelines to avoid reloading ---
+loaded_pipelines = {}
 
-    headers = {"Authorization": f"Bearer {token}"}
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "width": int(width),
-            "height": int(height),
-            "guidance_scale": float(guidance_scale),
-            "num_inference_steps": int(steps)
-        },
-        "options": {"wait_for_model": True}
-    }
+def get_pipeline(model_id):
+    """Load pipeline if not cached"""
+    if model_id not in loaded_pipelines:
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            use_safetensors=True
+        )
+        pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+        loaded_pipelines[model_id] = pipe
+    return loaded_pipelines[model_id]
 
+# --- Image generation function ---
+def generate_image(prompt, model_choice, width, height, guidance_scale, steps):
     try:
-        response = requests.post(HF_API_URL, headers=headers, json=payload, timeout=60)
-        response.raise_for_status()
+        model_id = MODEL_CHOICES[model_choice]
+        pipe = get_pipeline(model_id)
+
+        image = pipe(
+            prompt,
+            width=int(width),
+            height=int(height),
+            guidance_scale=float(guidance_scale),
+            num_inference_steps=int(steps)
+        ).images[0]
+
+        return image, f"✅ Generated with {model_choice}"
 
-        image = Image.open(io.BytesIO(response.content))
-        return image, "✅ Image generated successfully!"
     except Exception as e:
         return None, f"⚠️ Error: {str(e)}"
 
@@ -52,7 +61,11 @@ with gr.Blocks(css=css, title="Stable Diffusion Text-to-Image") as demo:
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", placeholder="A futuristic city at night", lines=3)
-            hf_token = gr.Textbox(label="Hugging Face API Token", placeholder="Enter your HF token", type="password")
+            model_choice = gr.Dropdown(
+                list(MODEL_CHOICES.keys()),
+                value="Dreamlike Diffusion 1.0",
+                label="Choose Model"
+            )
             width = gr.Dropdown([256, 384, 512, 768, 1024], value=512, label="Width")
             height = gr.Dropdown([256, 384, 512, 768, 1024], value=512, label="Height")
             guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.1, label="Guidance Scale")
@@ -60,14 +73,14 @@ with gr.Blocks(css=css, title="Stable Diffusion Text-to-Image") as demo:
             generate_btn = gr.Button("Generate Image", variant="primary")
 
         with gr.Column():
-            output_image = gr.Image(label="Generated Image")
+            output_image = gr.Image(label="Generated Image", type="pil")
             status = gr.Textbox(label="Status", interactive=False)
 
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, hf_token, width, height, guidance, steps],
+        inputs=[prompt, model_choice, width, height, guidance, steps],
         outputs=[output_image, status]
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", share=False)
+    demo.launch(server_name="0.0.0.0")