Update app.py

app.py

@@ -3,17 +3,22 @@ import torch
 from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
 from PIL import Image
 
-# Load models with CPU-specific optimizations
+# 1. LOAD MODEL (Using bfloat16 to fit in 16GB CPU RAM)
 model_id = "Tongyi-MAI/Z-Image-Turbo"
 
-# Initialize both pipelines (they share the same components to save RAM)
+print("Loading model... this takes a minute on CPU.")
+# We use bfloat16 and low_cpu_mem_usage to prevent OOM crashes
 txt2img = AutoPipelineForText2Image.from_pretrained(
-    model_id, torch_dtype=torch.float32, use_safetensors=True
+    model_id, 
+    torch_dtype=torch.bfloat16, 
+    low_cpu_mem_usage=True
 )
 txt2img.to("cpu")
 
+# Reuse components for the Edit (i2i) pipeline to save memory
 img2img = AutoPipelineForImage2Image.from_pipe(txt2img)
 
+# --- FUNCTIONS ---
 def generate_t2i(prompt, steps, guidance):
     return txt2img(
         prompt=prompt, 
@@ -22,7 +27,9 @@ def generate_t2i(prompt, steps, guidance):
     ).images[0]
 
 def generate_i2i(image, prompt, strength, steps, guidance):
-    # Resize to prevent CPU OOM (Out of Memory)
+    if image is None:
+        return None
+    # Resize to 512 to keep the CPU from exploding
     image = image.convert("RGB").resize((512, 512))
     return img2img(
         prompt=prompt,
@@ -32,17 +39,18 @@ def generate_i2i(image, prompt, strength, steps, guidance):
         guidance_scale=guidance
     ).images[0]
 
-# --- UI Layout ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# --- UI LAYOUT ---
+with gr.Blocks() as demo:
     gr.Markdown("# 🚀 Z-Image Turbo (CPU Edition)")
+    gr.Markdown("Note: CPU generation is slow (1-3 mins). Use 512x512 for best results.")
     
     with gr.Tabs():
         # Text to Image Tab
         with gr.TabItem("Text to Image"):
             with gr.Row():
                 with gr.Column():
-                    t2i_prompt = gr.Textbox(label="Prompt", placeholder="A dark fantasy castle...")
-                    t2i_steps = gr.Slider(1, 12, value=8, step=1, label="Steps (Turbo likes 8-10)")
+                    t2i_prompt = gr.Textbox(label="Prompt", placeholder="A liminal beach at night...")
+                    t2i_steps = gr.Slider(1, 12, value=8, step=1, label="Steps")
                     t2i_guidance = gr.Slider(0.0, 2.0, value=0.0, step=0.1, label="Guidance Scale")
                     t2i_btn = gr.Button("Generate")
                 with gr.Column():
@@ -50,17 +58,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             t2i_btn.click(generate_t2i, inputs=[t2i_prompt, t2i_steps, t2i_guidance], outputs=t2i_output)
 
         # Image to Image Tab
-        with gr.TabItem("Image to Image (The Editor)"):
+        with gr.TabItem("Image to Image"):
             with gr.Row():
                 with gr.Column():
-                    i2i_input = gr.Image(type="pil", label="Upload Source Image")
-                    i2i_prompt = gr.Textbox(label="Edit Prompt", placeholder="Change the style to...")
-                    i2i_strength = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Edit Strength (Lower = closer to original)")
+                    i2i_input = gr.Image(type="pil", label="Upload Photo")
+                    i2i_prompt = gr.Textbox(label="Edit Prompt", placeholder="Make it look like a dark fantasy painting")
+                    i2i_strength = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Edit Strength")
                     i2i_steps = gr.Slider(1, 12, value=10, step=1, label="Steps")
                     i2i_guidance = gr.Slider(0.0, 2.0, value=0.0, step=0.1, label="Guidance Scale")
-                    i2i_btn = gr.Button("Transform Image")
+                    i2i_btn = gr.Button("Apply Edits")
                 with gr.Column():
                     i2i_output = gr.Image(label="Result")
             i2i_btn.click(generate_i2i, inputs=[i2i_input, i2i_prompt, i2i_strength, i2i_steps, i2i_guidance], outputs=i2i_output)
 
-demo.launch()
+# CRITICAL: theme moved here, ssr_mode=False added to fix the -1 error
+demo.launch(theme=gr.themes.Soft(), ssr_mode=False)