Update app.py

app.py

@@ -1,67 +1,22 @@
 import os
-
-# ---- HARD CPU SAFETY ----
-os.environ["DIFFUSERS_NO_BITSANDBYTES"] = "1"
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-
 import torch
-import gradio as gr
-from diffusers import Flux2KleinPipeline
+from diffusers import FluxPipeline
 
-DEVICE = "cpu"
-DTYPE = torch.float32
-MAX_SIZE = 512
-
-print("Loading FLUX.2 Klein 4B (CPU mode)...")
+# Force CPU
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["DIFFUSERS_NO_BITSANDBYTES"] = "1"
 
-pipe = Flux2KleinPipeline.from_pretrained(
-    "black-forest-labs/FLUX.2-klein-4B",
-    torch_dtype=DTYPE,
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.float32
 )
-pipe.to(DEVICE)
-
-pipe.enable_attention_slicing()
-
-def generate(
-    prompt,
-    steps=2,
-    guidance=1.0,
-    width=384,
-    height=384,
-):
-    image = pipe(
-        prompt=prompt,
-        num_inference_steps=steps,
-        guidance_scale=guidance,
-        width=width,
-        height=height,
-    ).images[0]
-    return image
-
-with gr.Blocks() as demo:
-    gr.Markdown("# FLUX.2 Klein 4B — CPU (Experimental)")
-
-    prompt = gr.Textbox(
-        label="Prompt",
-        placeholder="A watercolor painting of a small house in the mountains",
-    )
-
-    with gr.Row():
-        steps = gr.Slider(1, 4, value=2, step=1, label="Steps")
-        guidance = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="CFG")
-
-    with gr.Row():
-        width = gr.Slider(256, MAX_SIZE, value=384, step=64, label="Width")
-        height = gr.Slider(256, MAX_SIZE, value=384, step=64, label="Height")
 
-    run = gr.Button("Generate (slow)")
-    output = gr.Image()
+pipe.to("cpu")
 
-    run.click(
-        fn=generate,
-        inputs=[prompt, steps, guidance, width, height],
-        outputs=output,
-    )
+image = pipe(
+    prompt="a cinematic portrait of a cyberpunk warrior",
+    num_inference_steps=20
+).images[0]
 
-demo.launch()
+image.save("output.png")
+print("✅ Generated on CPU")