using runwayml/stable-diffusion-v1-5

app.py

@@ -1,42 +1,49 @@
 import gradio as gr
 import torch
-from huggingface_hub import login
+from diffusers import OnnxStableDiffusionPipeline
 import os
 
 # Use CPU optimized model
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 
-# FLUX model requires authentication
-MODEL_ID = "black-forest-labs/FLUX.1-dev"
+# Using a proper ONNX model for CPU inference
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
 
-# You'll need to add your Hugging Face token as a secret in your Space
-# Then uncomment the following lines:
-from huggingface_hub import login
-login(token=os.getenv("HF_TOKEN"))  # Add HF_TOKEN as a secret in your Space settings
-
-print(f"Loading model {MODEL_ID}...")
+print(f"Loading ONNX model {MODEL_ID}...")
 
 try:
-    from diffusers import FluxPipeline
-    
-    pipe = FluxPipeline.from_pretrained(
+    # Use ONNX pipeline for better CPU performance
+    pipe = OnnxStableDiffusionPipeline.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.bfloat16
+        revision="onnx",
+        provider="CPUExecutionProvider",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
     )
-    pipe = pipe.to("cpu")  # For CPU, but this will be very slow
+    pipe.safety_checker = None  # Disable safety checker for simplicity
     pipe.set_progress_bar_config(disable=True)
-    print("FLUX Model loaded successfully!")
+    print("ONNX Model loaded successfully!")
 except Exception as e:
-    print(f"Error loading FLUX model: {e}")
-    print("Please ensure you have access to the model and have set your HF_TOKEN")
-    raise e
+    print(f"Failed to load ONNX model: {e}")
+    print("Falling back to regular Stable Diffusion pipeline...")
+    from diffusers import StableDiffusionPipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
+    )
+    pipe = pipe.to("cpu")
+    pipe.safety_checker = None
+    pipe.set_progress_bar_config(disable=True)
+    print("Model loaded!")
 
 def generate(prompt):
     try:
         result = pipe(
             prompt=prompt,
-            num_inference_steps=20,
-            guidance_scale=3.5  # FLUX typically uses lower guidance scale
+            num_inference_steps=40,
+            num_images_per_prompt=1,
+            guidance_scale=7.5
         )
         return result.images[0]
     except Exception as e:
@@ -47,8 +54,8 @@ interface = gr.Interface(
     fn=generate,
     inputs=gr.Textbox(lines=3, label="Prompt", placeholder="Enter your image prompt here..."),
     outputs=gr.Image(type="pil", label="Generated Image"),
-    title="⚡ FLUX Image Generation",
-    description="Image generation with FLUX model (requires authentication)"
+    title="⚡ ONNX Stable Diffusion",
+    description="Faster CPU-based image generation with ONNX optimization"
 )
 
 interface.launch()