Update app.py

app.py

@@ -1,27 +1,47 @@
-import gradio as gr
-import requests
-import io
 import os
-from PIL import Image
+import torch
+import gradio as gr
+from diffusers import StableDiffusion3Pipeline
+from huggingface_hub import login
+
+# Load Hugging Face token from environment variable
+hf_token = os.getenv("HF_TOKEN")
+
+# Authenticate using the Hugging Face token
+if hf_token is not None:
+    login(token=hf_token)
+else:
+    raise ValueError("Please set your Hugging Face token as the HF_TOKEN environment variable.")
+
+def image_generator(prompt):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Load the model using the authenticated token
+    pipeline = StableDiffusion3Pipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium-diffusers",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        text_encoder_3=None,
+        tokenizer_3=None,
+        use_auth_token=hf_token  # Use the token for authentication
+    )
+    pipeline.to(device)
 
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-headers = {"Authorization": f"Bearer {os.getenv('HF_AUTH_TOKEN')}"}
+    image = pipeline(
+        prompt=prompt,
+        negative_prompt="blurred, ugly, watermark, low, resolution, blurry",
+        num_inference_steps=40,
+        height=1024,
+        width=1024,
+        guidance_scale=9.0
+    ).images[0]
 
-def query_image(prompt):
-    payload = {"inputs": prompt}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    image_bytes = response.content
-    image = Image.open(io.BytesIO(image_bytes))
     return image
 
-# Gradio interface
-iface = gr.Interface(
-    fn=query_image,
-    inputs=gr.Textbox(label="Enter Prompt", placeholder="Give me lodhrangpt image"),
-    outputs=gr.Image(label="Generated Image"),
-    title="Image Generator",
-    description="Enter a prompt to generate an image using the FLUX model."
+interface = gr.Interface(
+    fn=image_generator,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."),
+    outputs=gr.Image(type="pil"),
+    title="Image Generator App",
+    description="This is a simple image generator app using HuggingFace's Stable Diffusion 3 model."
 )
 
-# Launch the Gradio app
-iface.launch()
+interface.launch()