Update app.py

app.py

@@ -1,38 +1,49 @@
+import os
 import torch
 import gradio as gr
 from diffusers import StableDiffusion3Pipeline
+from huggingface_hub import login
 
-def image_generator(prompt):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    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)
-    #pipeline.enable_model_cpu_offload()
-    pipeline.to(device)
-
-    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]
-
-    #image.show()
-    return image
+# Load Hugging Face token from environment variable
+hf_token = os.getenv("HF_TOKEN")
 
-#image_generator("A magical cat doing spell")
+login(token=hf_token)
 
-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.")
+# Set device (CUDA if available)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load the model once at the beginning
+pipeline = StableDiffusion3Pipeline.from_pretrained(
+    "stabilityai/stable-diffusion-3-medium-diffusers",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    use_auth_token=hf_token,  # Use the token for authentication
+    # Load the model in 8-bit precision
+    load_in_8bit=True
+)
+pipeline.to(device)
 
-interface.launch()
+def image_generator(prompt):
+    # Generate the image based on the prompt
+    with torch.no_grad():  # Disable gradient calculations for faster inference
+        image = pipeline(
+            prompt=prompt,
+            negative_prompt="blurred, ugly, watermark, low, resolution, blurry",
+            num_inference_steps=20,  # Further reduced steps for faster generation
+            height=512,  # Smaller height for quicker generation
+            width=512,  # Smaller width for quicker generation
+            guidance_scale=7.5,  # Moderate guidance scale
+            num_images_per_prompt=1  # Generate only one image per prompt
+        ).images[0]
 
+    return image
 
-print(interface)
+# Gradio interface setup
+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 interface
+interface.launch()