Update app.py

app.py

@@ -1,65 +1,37 @@
-import os
 import gradio as gr
-import torch
-import spaces
-from diffusers import FluxPipeline
-from huggingface_hub import login
+from optimum.intel import OVStableDiffusionPipeline
 
-# --- AUTHENTICATION ---
-token = os.getenv("HF_TOKEN")
-if token:
-    login(token=token)
+# 1. Load the OpenVINO optimized model
+# This model is specifically 'squeezed' to run on CPUs
+model_id = "OpenVINO/stable-diffusion-v1-5-int8-ov"
 
-# --- MODEL LOADING ---
-model_id = "black-forest-labs/FLUX.1-schnell"
+print("Loading model to CPU... (This may take a minute)")
+pipe = OVStableDiffusionPipeline.from_pretrained(model_id, compile=True)
 
-# Load without .to("cuda") here. Let ZeroGPU handle it.
-pipe = FluxPipeline.from_pretrained(
-    model_id, 
-    torch_dtype=torch.bfloat16,
-    token=token
-)
+# 2. Define the Generation Function
+def generate_image(prompt, steps):
+    # On CPU, we keep steps low (1-4) for speed
+    image = pipe(
+        prompt=prompt, 
+        num_inference_steps=steps,
+        guidance_scale=7.5
+    ).images[0]
+    return image
 
-@spaces.GPU(duration=60)
-def generate_image(prompt):
-    if not prompt:
-        return None
+# 3. Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ⚡ CPU-Fast Image Generator")
+    gr.Markdown("Using OpenVINO to run Stable Diffusion on a free CPU tier.")
     
-    try:
-        # ZeroGPU specific: ensure the pipe is on the right device
-        pipe.to("cuda") 
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Enter Prompt", value="A cozy cabin in the woods, digital art")
+            steps = gr.Slider(1, 10, value=4, step=1, label="Steps (Lower is faster)")
+            btn = gr.Button("Generate", variant="primary")
         
-        generator = torch.Generator("cuda").manual_seed(42)
-            
-        image = pipe(
-            prompt=prompt,
-            width=1024,
-            height=1024,
-            num_inference_steps=4, 
-            generator=generator,
-            guidance_scale=0.0
-        ).images[0]
-            
-        return image
-    except Exception as e:
-        # This will show the actual error message in the Gradio UI
-        raise gr.Error(f"Generation failed: {str(e)}")
+        output_img = gr.Image(label="Result")
 
-# --- MINIMALIST UI ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# FLUX.1 Minimal")
-    
-    with gr.Group():
-        prompt = gr.Textbox(
-            placeholder="Describe your image...",
-            show_label=False,
-            container=False
-        )
-        btn = gr.Button("Generate", variant="primary")
-    
-    output_img = gr.Image(show_label=False, interactive=False)
-    
-    btn.click(fn=generate_image, inputs=[prompt], outputs=output_img)
+    btn.click(fn=generate_image, inputs=[prompt, steps], outputs=output_img)
 
 if __name__ == "__main__":
     demo.launch(share=True)