Update app.py

app.py

@@ -1,26 +1,52 @@
-from diffusers import StableDiffusionXLPipeline
-import torch
+import requests
+import io
+from PIL import Image
 import gradio as gr
+import os  # To access environment variables
 
-# Load the Stable Diffusion XL model
-pipeline = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
-).to("cuda")
+# Access the Hugging Face API token securely
+API_TOKEN = os.getenv("HF_API_TOKEN")
 
-def generate_image(prompt, prompt_2):
-    # Generate the image based on the prompts
-    image = pipeline(prompt=prompt, prompt_2=prompt_2).images[0]
-    
-    # Convert the PIL image to a format that Gradio can display
+if not API_TOKEN:
+    raise ValueError("Hugging Face API token not found. Please check your Space's secrets configuration.")
+
+API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+
+def generate_image(prompt, negative_prompt, guidance_scale, width, height, num_inference_steps):
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "negative_prompt": negative_prompt,
+            "guidance_scale": guidance_scale,
+            "width": width,
+            "height": height,
+            "num_inference_steps": num_inference_steps,
+        },
+    }
+    image_bytes = query(payload)
+    image = Image.open(io.BytesIO(image_bytes))
     return image
 
-# Define the Gradio interface
-iface = gr.Interface(fn=generate_image,
-                     inputs=[gr.Textbox(label="Prompt 1: Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"),
-                             gr.Textbox(label="Prompt 2: Van Gogh painting")],
-                     outputs="image",
-                     title="Stable Diffusion XL Image Generator",
-                     description="Generate images with Stable Diffusion XL based on two prompts.")
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Negative Prompt"),
+        gr.Slider(label="Guidance Scale", minimum=1, maximum=20, step=0.1, default=7.5),
+        gr.Slider(label="Width", minimum=768, maximum=1024, step=1, default=1024),
+        gr.Slider(label="Height", minimum=768, maximum=1024, step=1, default=768),
+        gr.Slider(label="Number of Inference Steps", minimum=20, maximum=50, step=1, default=30)
+    ],
+    outputs=gr.Image(type="pil"),
+    title="Stable Diffusion XL Image Generator",
+    description="Generate images with Stable Diffusion XL. Provide a prompt, specify any negative prompts, and adjust the image generation parameters.",
+)
 
-# Launch the Gradio app
-iface.launch()
+# Launch the Gradio app, setting share=True for Hugging Face Spaces
+iface.launch(share=True)