Update app.py

app.py

@@ -1,22 +1,13 @@
 import gradio as gr
-import torch
-from diffusers import StableDiffusionPipeline
 import random
 import numpy as np
+import os
+from huggingface_hub import InferenceClient
 
-# 1. Load a CPU-optimized model
-# 'segmind/tiny-sd' is much smaller and faster on CPUs than SDXL
-model_id = "segmind/tiny-sd"
-
-# Use float32 because CPU doesn't support float16 well
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float32
-)
-pipe = pipe.to("cpu")
-
-# Optimize for CPU speed
-pipe.set_progress_bar_config(disable=True)
+# 1. Initialize the Client
+# If you are running this IN a Hugging Face Space, it will automatically
+# find your token if you add it to 'Settings > Variables and Secrets' as HF_TOKEN
+client = InferenceClient("black-forest-labs/FLUX.1-schnell")
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -24,33 +15,43 @@ def infer(prompt, seed, randomize_seed, width, height):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
-    generator = torch.Generator("cpu").manual_seed(seed)
-    
-    # We use very low steps (10-15) because CPU is slow
-    image = pipe(
-        prompt=prompt,
-        generator=generator,
-        num_inference_steps=15, 
-        guidance_scale=7.0,
-        width=width,
-        height=height
-    ).images[0]
-    
-    return image, seed
+    # FLUX.1-schnell is optimized for 4 steps. 
+    # The Inference Client handles the generator/torch logic server-side.
+    try:
+        image = client.text_to_image(
+            prompt,
+            width=width,
+            height=height,
+            num_inference_steps=4, # Schnell specifically needs only 4 steps
+            guidance_scale=0.0,    # Schnell usually ignores guidance or prefers 0.0
+            seed=seed
+        )
+        return image, seed
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {e}")
 
 # Simple UI
-with gr.Blocks() as demo:
-    gr.Markdown("# CodeIgnite CPU Image Engine")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚀 CodeIgnite FLUX Engine")
+    gr.Markdown("Using `FLUX.1-schnell` via Inference API for lightning-fast results.")
+    
     with gr.Column():
-        prompt = gr.Textbox(label="Prompt", placeholder="A simple cat drawing")
-        run_button = gr.Button("Generate (CPU Mode)")
+        with gr.Row():
+            prompt = gr.Textbox(
+                label="Prompt", 
+                placeholder="A futuristic cyberpunk city...", 
+                scale=4
+            )
+            run_button = gr.Button("Generate", variant="primary")
+        
         result = gr.Image(label="Result")
         
-        with gr.Accordion("Settings", open=False):
+        with gr.Accordion("Advanced Settings", open=False):
             seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            width = gr.Slider(label="Width", minimum=256, maximum=512, step=32, value=384)
-            height = gr.Slider(label="Height", minimum=256, maximum=512, step=32, value=384)
+            # FLUX works best at 1024x1024, but 512-1024 is safe for API
+            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=32, value=1024)
+            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=32, value=1024)
 
     run_button.click(
         fn=infer,