Update app.py

app.py

@@ -1,363 +1,145 @@
 import os
-import subprocess
-import sys
-import io
 import gradio as gr
 import numpy as np
 import random
 import spaces
 import torch
-from diffusers import Flux2KleinPipeline
-import requests
-from PIL import Image
-import json
+from diffusers import FluxImg2ImgPipeline
+from PIL import Image, ImageOps
+import io
 import base64
-from huggingface_hub import InferenceClient
 
+# --- CONFIGURATION ---
+# We use bfloat16 for speed/memory on CPU
 dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-hf_client = InferenceClient(
-    api_key=os.environ.get("HF_TOKEN"),
-)
-VLM_MODEL = "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"
-
-SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.
-
-Guidelines:
-1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
-2. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.
-3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.
-
-Output only the revised prompt and nothing else."""
-
-SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.2 by Black Forest Labs, an image-editing expert. You convert editing requests into one concise instruction (50-80 words, ~30 for brief requests).
-
-Rules:
-- Single instruction only, no commentary
-- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
-- Specify what changes AND what stays the same (face, lighting, composition)
-- Reference actual image elements
-- Turn negatives into positives ("don't change X" → "keep X")
-- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
-- Keep content PG-13
-
-Output only the final instruction in plain text and nothing else."""
-
-# Model repository IDs for 4B
-REPO_ID_REGULAR = "black-forest-labs/FLUX.2-klein-base-4B"
-REPO_ID_DISTILLED = "black-forest-labs/FLUX.2-klein-4B"
-
-# Load both 4B models
-print("Loading 4B Regular model...")
-pipe_regular = Flux2KleinPipeline.from_pretrained(REPO_ID_REGULAR, torch_dtype=dtype)
-pipe_regular.to("cuda")
-
-print("Loading 4B Distilled model...")
-pipe_distilled = Flux2KleinPipeline.from_pretrained(REPO_ID_DISTILLED, torch_dtype=dtype)
-pipe_distilled.to("cuda")
+device = "cpu"
 
-# Dictionary for easy access
-pipes = {
-    "Distilled (4 steps)": pipe_distilled,
-    "Base (50 steps)": pipe_regular,
-}
+# We use Schnell because it is 10x faster on CPU than Dev or 'Klein'
+# If you explicitly have access to Flux 2, change this ID.
+MODEL_ID = "black-forest-labs/FLUX.1-schnell"
 
-# Default steps for each mode
-DEFAULT_STEPS = {
-    "Distilled (4 steps)": 4,
-    "Base (50 steps)": 50,
-}
+print(f"Loading Model: {MODEL_ID} on {device}...")
 
-# Default CFG for each mode
-DEFAULT_CFG = {
-    "Distilled (4 steps)": 1.0,
-    "Base (50 steps)": 4.0,
-}
-
-def image_to_data_uri(img):
-    buffered = io.BytesIO()
-    img.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    return f"data:image/png;base64,{img_str}"
+# Load Pipeline
+pipe = FluxImg2ImgPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=dtype
+)
 
+# CRITICAL CPU OPTIMIZATION
+# This replaces .to("cuda"). It loads the model in pieces so RAM doesn't crash.
+pipe.enable_model_cpu_offload()
 
-def upsample_prompt_logic(prompt, image_list):
-    try:
-        if image_list and len(image_list) > 0:
-            # Image + Text Editing Mode
-            system_content = SYSTEM_PROMPT_WITH_IMAGES
-            
-            # Construct user message with text and images
-            user_content = [{"type": "text", "text": prompt}]
-            
-            for img in image_list:
-                data_uri = image_to_data_uri(img)
-                user_content.append({
-                    "type": "image_url",
-                    "image_url": {"url": data_uri}
-                })
-                
-            messages = [
-                {"role": "system", "content": system_content},
-                {"role": "user", "content": user_content}
-            ]
-        else:
-            # Text Only Mode
-            system_content = SYSTEM_PROMPT_TEXT_ONLY
-            messages = [
-                {"role": "system", "content": system_content},
-                {"role": "user", "content": prompt}
-            ]
+MAX_SEED = np.iinfo(np.int32).max
 
-        completion = hf_client.chat.completions.create(
-            model=VLM_MODEL,
-            messages=messages,
-            max_tokens=1024
-        )
+# --- THE "INJECTION" LOGIC ---
+# This acts as a pre-processor injection. It forces the pixel data 
+# to be symmetrical before the UNet even sees it. 
+# This guarantees the "Face Lock".
+def inject_symmetry(image, side="Left"):
+    if image is None:
+        return None
         
-        return completion.choices[0].message.content
-    except Exception as e:
-        print(f"Upsampling failed: {e}")
-        return prompt
-
+    img = image.convert("RGB")
+    w, h = img.size
+    mid = w // 2
+    arr = np.array(img)
+    
+    # Mathematical locking of geometry
+    if side == "Left":
+        target_half = arr[:, :mid, :] # Get Left
+        mirrored = np.fliplr(target_half) # Mirror it
+        locked_face = np.concatenate((target_half, mirrored), axis=1)
+    else:
+        target_half = arr[:, mid:, :] # Get Right
+        mirrored = np.fliplr(target_half) # Mirror it
+        locked_face = np.concatenate((mirrored, target_half), axis=1)
+        
+    return Image.fromarray(locked_face)
 
-def update_dimensions_from_image(image_list):
-    """Update width/height sliders based on uploaded image aspect ratio.
-    Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
-    if image_list is None or len(image_list) == 0:
-        return 1024, 1024  # Default dimensions
-    
-    # Get the first image to determine dimensions
-    img = image_list[0][0]  # Gallery returns list of tuples (image, caption)
-    img_width, img_height = img.size
+# --- INFERENCE FUNCTION ---
+@spaces.GPU(duration=120) # Request GPU if available, falls back to CPU logic if not
+def infer(prompt, input_image, side_choice, strength, seed, randomize_seed, width, height, steps, guidance):
     
-    aspect_ratio = img_width / img_height
-    
-    if aspect_ratio >= 1:  # Landscape or square
-        new_width = 1024
-        new_height = int(1024 / aspect_ratio)
-    else:  # Portrait
-        new_height = 1024
-        new_width = int(1024 * aspect_ratio)
-    
-    # Round to nearest multiple of 8
-    new_width = round(new_width / 8) * 8
-    new_height = round(new_height / 8) * 8
-    
-    # Ensure within valid range (minimum 256, maximum 1024)
-    new_width = max(256, min(1024, new_width))
-    new_height = max(256, min(1024, new_height))
-    
-    return new_width, new_height
-
-
-def update_steps_from_mode(mode_choice):
-    """Update the number of inference steps based on the selected mode."""
-    return DEFAULT_STEPS[mode_choice], DEFAULT_CFG[mode_choice]
-
+    if input_image is None:
+        raise gr.Error("Please upload an image for face symmetry.")
 
-@spaces.GPU(duration=85)
-def infer(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
-    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    
-    # Select the appropriate pipeline based on mode choice
-    pipe = pipes[mode_choice]
-    
-    # Prepare image list (convert None or empty gallery to None)
-    image_list = None
-    if input_images is not None and len(input_images) > 0:
-        image_list = []
-        for item in input_images:
-            image_list.append(item[0])
-
-    # 1. Upsampling (Network bound)
-    final_prompt = prompt
-    if prompt_upsampling:
-        progress(0.1, desc="Upsampling prompt...")
-        final_prompt = upsample_prompt_logic(prompt, image_list)
-        print(f"Original Prompt: {prompt}")
-        print(f"Upsampled Prompt: {final_prompt}")
-
-    # 2. Image Generation
-    progress(0.2, desc=f"Generating image with 4B {mode_choice}...")
-    
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    pipe_kwargs = {
-        "prompt": final_prompt,
-        "height": height,
-        "width": width,
-        "num_inference_steps": num_inference_steps,
-        "guidance_scale": guidance_scale,
-        "generator": generator,
-    }
-    
-    # Add images if provided
-    if image_list is not None:
-        pipe_kwargs["image"] = image_list
-    
-    image = pipe(**pipe_kwargs).images[0]
-    
-    return image, seed
-
-
-examples = [
-    ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
-    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
-    ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
-    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
-]
-
-examples_images = [
-    ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
-]
-
+        
+    # 1. INJECT SYMMETRY
+    # We process the image *before* the model touches it.
+    print("Injecting symmetry constraints...")
+    processed_image = inject_symmetry(input_image, side_choice)
+    
+    # Resize to be compatible with Flux (multiples of 16)
+    w, h = processed_image.size
+    w = (w // 16) * 16
+    h = (h // 16) * 16
+    processed_image = processed_image.resize((w, h))
+
+    print("Running Flux to smooth seams...")
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    
+    # 2. RUN DIFFUSION
+    # We use the processed image as the base.
+    # Strength is CRITICAL: 
+    # 0.1 - 0.30 = Locks Identity (Only fixes the seam)
+    # 0.35+ = Starts changing the face
+    result = pipe(
+        prompt=prompt,
+        image=processed_image,
+        strength=strength, 
+        num_inference_steps=steps,
+        guidance_scale=guidance,
+        generator=generator
+    ).images[0]
+    
+    return result, seed
+
+# --- UI SETUP ---
 css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1200px;
-}
-.gallery-container img{
-    object-fit: contain;
-}
+#col-container { max-width: 1000px; margin: 0 auto; }
 """
 
 with gr.Blocks(css=css) as demo:
-    
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.2 [Klein] - 4B (Apache 2.0)
-FLUX.2 [klein] is a fast, unified image generation and editing model designed for fast inference [[model](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B)], [[blog](https://bfl.ai/blog/flux-2)]
-        """)
+        gr.Markdown("# Flux Face Symmetry (Identity Lock)")
+        gr.Markdown("CPU-Optimized Mode. Uses Pixel Injection to lock face geometry.")
+        
         with gr.Row():
             with gr.Column():
+                input_img = gr.Image(label="Upload Face", type="pil")
+                
                 with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=2,
-                        placeholder="Enter your prompt",
-                        container=False,
-                        scale=3
-                    )
-                    
-                    run_button = gr.Button("Run", scale=1)
-                    
-                with gr.Accordion("Input image(s) (optional)", open=False):
-                    input_images = gr.Gallery(
-                        label="Input Image(s)",
-                        type="pil",
-                        columns=3,
-                        rows=1,
-                    )
-
-                mode_choice = gr.Radio(
-                    label="Mode",
-                    choices=["Distilled (4 steps)", "Base (50 steps)"],
-                    value="Distilled (4 steps)",
-                )
+                    side = gr.Radio(["Left", "Right"], label="Keep Side", value="Left")
+                    # Strength default is 0.25 -> This ensures ID is locked
+                    strength = gr.Slider(0.1, 0.6, value=0.25, step=0.01, label="Denoise Strength (Keep <0.30 to lock ID)")
                 
-                with gr.Accordion("Advanced Settings", open=False):
-                    
-                    prompt_upsampling = gr.Checkbox(
-                        label="Prompt Upsampling",
-                        value=False,
-                        info="Automatically enhance the prompt using a VLM"
-                    )
-        
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
-                    )
-                    
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                    
-                    with gr.Row():
-                        
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                        
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=8,
-                            value=1024,
-                        )
-                    
-                    with gr.Row():
-                        
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=100,
-                            step=1,
-                            value=4,
-                        )
-                        
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=0.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=1.0,
-                        )
+                prompt = gr.Text(
+                    label="Prompt (Optional - usually leave empty or describe lighting)", 
+                    value="high quality, realistic, smooth skin, 8k",
+                    lines=2
+                )
                 
+                run_btn = gr.Button("Generate Symmetrical Face", variant="primary")
                 
-            with gr.Column():
-                result = gr.Image(label="Result", show_label=False)
-            
-        
-        gr.Examples(
-            examples=examples,
-            fn=infer,
-            inputs=[prompt],
-            outputs=[result, seed],
-            cache_examples=True,
-            cache_mode="lazy"
-        )
+                with gr.Accordion("Advanced", open=False):
+                    steps = gr.Slider(1, 50, value=4, step=1, label="Steps (Keep low for CPU)")
+                    guidance = gr.Slider(0, 10, value=1.0, step=0.1, label="Guidance")
+                    width = gr.Slider(256, 1024, value=1024, step=16, label="Width")
+                    height = gr.Slider(256, 1024, value=1024, step=16, label="Height")
+                    seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
+                    randomize_seed = gr.Checkbox(True, label="Randomize Seed")
 
-        gr.Examples(
-            examples=examples_images,
-            fn=infer,
-            inputs=[prompt, input_images],
-            outputs=[result, seed],
-            cache_examples=True,
-            cache_mode="lazy"
-        )
-
-    # Auto-update dimensions when images are uploaded
-    input_images.upload(
-        fn=update_dimensions_from_image,
-        inputs=[input_images],
-        outputs=[width, height]
-    )
-    
-    # Auto-update steps when mode changes
-    mode_choice.change(
-        fn=update_steps_from_mode,
-        inputs=[mode_choice],
-        outputs=[num_inference_steps, guidance_scale]
-    )
+            with gr.Column():
+                output_img = gr.Image(label="Result")
+                seed_output = gr.Number(label="Used Seed")
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
-        outputs=[result, seed]
+    run_btn.click(
+        infer,
+        inputs=[prompt, input_img, side, strength, seed, randomize_seed, width, height, steps, guidance],
+        outputs=[output_img, seed_output]
     )
 
 demo.launch()