Upload app.py with huggingface_hub

app.py

@@ -5,14 +5,24 @@ import random
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-client_hf = InferenceClient(provider="hf-inference", token=HF_TOKEN)
-client_fal = InferenceClient(provider="fal-ai", api_key=HF_TOKEN)
+client = InferenceClient(provider="fal-ai", api_key=HF_TOKEN)
 
 MODELS = {
-    "SDXL Base": {"model_id": "stabilityai/stable-diffusion-xl-base-1.0", "client": client_fal},
-    "FLUX.1-schnell": {"model_id": "black-forest-labs/FLUX.1-schnell", "client": client_fal},
-    "FLUX.1-dev": {"model_id": "black-forest-labs/FLUX.1-dev", "client": client_fal},
-    "SD v1.5 (Free)": {"model_id": "stable-diffusion-v1-5/stable-diffusion-v1-5", "client": client_hf},
+    "SDXL Base": {
+        "model_id": "stabilityai/stable-diffusion-xl-base-1.0",
+        "max_steps": 50,
+        "recommended_steps": 30,
+    },
+    "FLUX.1-schnell (4 steps)": {
+        "model_id": "black-forest-labs/FLUX.1-schnell",
+        "max_steps": 4,
+        "recommended_steps": 4,
+    },
+    "FLUX.1-dev (28 steps)": {
+        "model_id": "black-forest-labs/FLUX.1-dev",
+        "max_steps": 28,
+        "recommended_steps": 28,
+    },
 }
 
 DEFAULT_MODEL = "SDXL Base"
@@ -40,41 +50,32 @@ def generate_image(prompt, model_name=DEFAULT_MODEL, style="None", negative_prom
         seed = random.randint(0, 2**32 - 1)
     model_info = MODELS.get(model_name, MODELS[DEFAULT_MODEL])
     model_id = model_info["model_id"]
-    client = model_info["client"]
+    max_steps = model_info["max_steps"]
+    
+    # Clamp inference steps to model's max supported
+    steps = min(num_inference_steps, max_steps)
 
     try:
         image = client.text_to_image(
             prompt=full_prompt, model=model_id,
             negative_prompt=negative_prompt if negative_prompt.strip() else None,
-            guidance_scale=guidance_scale, num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale, num_inference_steps=steps,
             width=width, height=height, seed=seed,
         )
-        return image, f"Seed: {seed} | Model: {model_name}"
+        return image, f"Seed: {seed} | Model: {model_name} | Steps: {steps}"
     except Exception as e:
-        error_msg = str(e)
-        if client != client_hf:
-            try:
-                image = client_hf.text_to_image(
-                    prompt=full_prompt, model="stable-diffusion-v1-5/stable-diffusion-v1-5",
-                    negative_prompt=negative_prompt if negative_prompt.strip() else None,
-                    guidance_scale=guidance_scale, num_inference_steps=num_inference_steps,
-                    width=min(width, 512), height=min(height, 512), seed=seed,
-                )
-                return image, f"Fallback SD v1.5 | Seed: {seed}" + f" | Size: {min(width, 512)}x{min(height, 512)}"
-            except Exception as e2:
-                raise gr.Error(f"Primary: {error_msg} | Fallback: {str(e2)}")
-        raise gr.Error(f"Generation failed: {error_msg}")
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 
 EXAMPLES = [
     ["A majestic dragon flying over a crystal lake at sunset, epic fantasy art", "SDXL Base", "Fantasy", "", 7.5, 30, -1, 1024, 1024],
     ["A cute robot exploring a colorful garden", "SDXL Base", "Anime", "", 7.5, 30, -1, 1024, 1024],
-    ["An astronaut riding a horse on Mars", "FLUX.1-schnell", "Photorealistic", "", 7.5, 30, -1, 1024, 1024],
+    ["An astronaut riding a horse on Mars", "FLUX.1-schnell (4 steps)", "Photorealistic", "", 7.5, 4, -1, 1024, 1024],
     ["A cozy coffee shop interior with rain outside", "SDXL Base", "Digital Art", "", 7.5, 30, -1, 1024, 1024],
     ["A futuristic city skyline at night", "SDXL Base", "3D Render", "", 7.5, 30, -1, 1024, 1024],
     ["Beautiful mountain landscape with a river, golden hour", "SDXL Base", "Oil Painting", "", 7.5, 30, -1, 1024, 1024],
-    ["A magical forest with glowing mushrooms", "FLUX.1-schnell", "Fantasy", "", 7.5, 30, -1, 1024, 1024],
-    ["Portrait of a wise old wizard", "FLUX.1-dev", "Digital Art", "", 7.5, 30, -1, 1024, 1024],
+    ["A magical forest with glowing mushrooms", "FLUX.1-schnell (4 steps)", "Fantasy", "", 7.5, 4, -1, 1024, 1024],
+    ["Portrait of a wise old wizard", "FLUX.1-dev (28 steps)", "Digital Art", "", 7.5, 28, -1, 1024, 1024],
 ]
 
 
@@ -85,7 +86,7 @@ with gr.Blocks(
     theme=gr.themes.Soft(),
     css=".generate-btn { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important; }"
 ) as demo:
-    gr.Markdown("# 🎨 Text-to-Image Generator\n\nGenerate images from text using SDXL, FLUX, and Stable Diffusion.")
+    gr.Markdown("# 🎨 Text-to-Image Generator\n\nGenerate images from text using SDXL and FLUX via HuggingFace Inference API.")
 
     with gr.Row():
         with gr.Column(scale=3):
@@ -140,7 +141,7 @@ with gr.Blocks(
         outputs=[output_image, info_text],
     )
 
-    gr.Markdown("---\n**Models:** SDXL (Stability AI) | FLUX (Black Forest Labs) | SD v1.5")
+    gr.Markdown("---\n**Models:** SDXL (Stability AI) | FLUX.1-schnell/dev (Black Forest Labs) via fal-ai")
 
 if __name__ == "__main__":
     demo.launch()