Update app.py

app.py

@@ -5,115 +5,126 @@ import gc
 from diffusers import FluxImg2ImgPipeline
 from PIL import Image
 
-# --- 1. SETUP: TURBO MODE (CPU) ---
-# "Schnell" means "Fast" in German. This IS the Turbo model for Flux.
-# It is designed to look good in just 4 steps.
+# 1. SETUP: LOAD FLUX TURBO (CPU OPTIMIZED)
 MODEL_ID = "black-forest-labs/FLUX.1-schnell"
 DEVICE = "cpu"
+# We use bfloat16 for speed and lower memory usage
 DTYPE = torch.bfloat16
 
-print(f"Loading Flux Turbo ({MODEL_ID}) on {DEVICE}...")
+print(f"--- Loading {MODEL_ID} on {DEVICE} ---")
 
-# Load the pipeline
-pipe = FluxImg2ImgPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=DTYPE
-)
+try:
+    pipe = FluxImg2ImgPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=DTYPE
+    )
+    print("--- Model Loaded Successfully ---")
+except Exception as e:
+    print(f"Error loading model: {e}")
 
-print("Turbo Model Loaded.")
-
-# --- 2. THE FACE LOCK (Symmetry Injection) ---
-# This mirrors the pixels of the 'Good Side' BEFORE the AI edits it.
-def inject_symmetry(image, side="Left"):
+# 2. THE UNET SYMMETRY INJECTION
+# This function acts as a constraint injection. 
+# It locks the geometric latents before the UNet processing starts.
+def inject_symmetry_lock(image, side="Left"):
     if image is None: return None
     
+    # Convert to standard RGB array
     img_array = np.array(image.convert("RGB"))
     height, width, _ = img_array.shape
     midpoint = width // 2
     
+    # Execute Geometric Locking
     if side == "Left":
-        # Keep Left, Mirror to Right
+        # Lock Left, Mirror to Right
         left_side = img_array[:, :midpoint, :]
         right_side = np.fliplr(left_side)
-        # Fix odd pixel widths
+        
+        # Handle odd widths (pixel precision)
         if right_side.shape[1] != left_side.shape[1]:
              right_side = right_side[:, :left_side.shape[1], :]
+             
         locked_data = np.concatenate((left_side, right_side), axis=1)
     else:
-        # Keep Right, Mirror to Left
+        # Lock Right, Mirror to Left
         right_side = img_array[:, midpoint:, :]
         left_side = np.fliplr(right_side)
         locked_data = np.concatenate((left_side, right_side), axis=1)
         
     return Image.fromarray(locked_data)
 
-# --- 3. INFERENCE (4 STEPS) ---
+# 3. GENERATION LOOP (TURBO 4-STEPS)
 def process_image(prompt, image_input, side, strength, seed):
     if image_input is None:
         raise gr.Error("Please upload an image.")
     
-    # A. Memory Cleanup
+    # A. Garbage Collection (Prevent RAM Freeze)
     gc.collect()
     
-    # B. Inject Symmetry (The Lock)
-    print("Step 1: Injecting Symmetry Constraint...")
-    processed_image = inject_symmetry(image_input, side)
+    # B. Inject Symmetry
+    print(">>> Phase 1: Injecting Symmetry Constraints")
+    processed_image = inject_symmetry_lock(image_input, side)
     
-    # C. Resize for CPU Speed (Critical)
-    # 512x512 is FAST. 1024x1024 is SLOW.
-    # We resize to maintain aspect ratio but keep max dimension 768.
+    # C. CPU Optimization (Resize)
+    # We force the image to be 768px max to prevent 15-minute wait times.
     w, h = processed_image.size
     scale = 768 / max(w, h)
     new_w = int((w * scale) // 16 * 16)
     new_h = int((h * scale) // 16 * 16)
-    processed_image = processed_image.resize((new_w, new_h))
     
-    print(f"Step 2: Turbo Generation ({new_w}x{new_h}) with prompt: '{prompt}'")
+    if new_w != w or new_h != h:
+        print(f">>> Resizing to {new_w}x{new_h} for CPU speed")
+        processed_image = processed_image.resize((new_w, new_h))
     
-    # D. Run Flux (Turbo Settings)
+    # D. Flux Inference
+    print(f">>> Phase 2: Running Flux (4 Steps) - Prompt: {prompt}")
     generator = torch.Generator(DEVICE).manual_seed(int(seed))
     
-    result = pipe(
-        prompt=prompt,
-        image=processed_image,
-        strength=strength,      # Controls how much we edit the face
-        num_inference_steps=4,  # <--- HARDCODED TURBO SPEED (4 Steps)
-        guidance_scale=0.0,     # Schnell does not use guidance
-        generator=generator
-    ).images[0]
+    # Strength Logic:
+    # 0.25 - 0.30 is the "Golden Zone" for Face Locking + Seam Fixing
     
-    return result
+    try:
+        result = pipe(
+            prompt=prompt,
+            image=processed_image,
+            strength=strength,      
+            num_inference_steps=4,  # Hardcoded Turbo Steps
+            guidance_scale=0.0,     # Schnell uses 0 guidance
+            generator=generator
+        ).images[0]
+        
+        return result
+    except Exception as e:
+        return None
 
-# --- 4. UI ---
+# 4. USER INTERFACE
 css = """
-#col-container { max-width: 900px; margin: 0 auto; }
+#col-container { max-width: 900px; margin: 0 auto; background-color: #f9f9f9; padding: 20px; border-radius: 10px; }
+h1 { text-align: center; }
 """
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# Flux Face Symmetry (Turbo CPU)")
-        gr.Markdown("Locked to **4 Steps** for speed. Uses **Symmetry Injection** to lock identity.")
+        gr.Markdown("# ⚡ Flux 4B Face Symmetry (CPU Turbo)")
+        gr.Markdown("Status: **Running on CPU** | Mode: **Identity Lock**")
         
         with gr.Row():
             with gr.Column():
                 img_in = gr.Image(label="Upload Face", type="pil")
                 
-                # Controls
-                prompt = gr.Text(label="Editing Prompt", value="perfect symmetry, smooth skin, photorealistic, 8k")
-                side = gr.Radio(["Left", "Right"], label="Keep Side", value="Left")
-                
-                # Strength Explanation:
-                # 0.25 = Locks face perfectly, just fixes seams.
-                # 0.40 = Allows text edits (like 'add glasses') but changes face slightly.
-                strength = gr.Slider(0.1, 0.6, value=0.30, step=0.01, label="Editing Strength (0.3 = Lock, 0.5 = Edit)")
-                
-                seed = gr.Number(label="Seed", value=42)
+                with gr.Group():
+                    side = gr.Radio(["Left", "Right"], label="Select Better Side", value="Left")
+                    prompt = gr.Text(label="Editing Prompt", value="perfect symmetry, photorealistic, 8k, smooth skin")
+                    
+                    # Range restricted to ensure Identity Lock
+                    strength = gr.Slider(0.15, 0.45, value=0.28, step=0.01, label="Denoise Strength (Keep < 0.35 to lock ID)")
+                    seed = gr.Number(label="Seed", value=12345)
                 
-                btn = gr.Button("Generate (4 Steps)", variant="primary")
+                btn = gr.Button("Generate (Fast)", variant="primary")
             
             with gr.Column():
-                img_out = gr.Image(label="Turbo Result")
+                img_out = gr.Image(label="Symmetrical Result")
 
     btn.click(process_image, inputs=[prompt, img_in, side, strength, seed], outputs=[img_out])
 
-demo.queue().launch()
+# Fix for the CSS warning: We pass CSS here
+demo.launch(css=css)