Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| import gc | |
| from diffusers import FluxImg2ImgPipeline | |
| from PIL import Image | |
| # 1. SETUP | |
| MODEL_ID = "black-forest-labs/FLUX.1-schnell" | |
| DEVICE = "cpu" | |
| DTYPE = torch.bfloat16 | |
| print(f"--- Loading {MODEL_ID} on {DEVICE} ---") | |
| pipe = FluxImg2ImgPipeline.from_pretrained(MODEL_ID, torch_dtype=DTYPE) | |
| print("--- Model Loaded ---") | |
| # 2. SYMMETRY LOGIC | |
| def inject_symmetry_lock(image, side="Left"): | |
| if image is None: return None | |
| img_array = np.array(image.convert("RGB")) | |
| height, width, _ = img_array.shape | |
| midpoint = width // 2 | |
| if side == "Left": | |
| left_side = img_array[:, :midpoint, :] | |
| right_side = np.fliplr(left_side) | |
| 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: | |
| 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. GENERATION (SAFE MODE) | |
| def process_image(prompt, image_input, side, strength, seed): | |
| if image_input is None: | |
| raise gr.Error("Please upload an image.") | |
| gc.collect() | |
| # Step A: Symmetry | |
| print(">>> Step 1: Injecting Symmetry") | |
| processed_image = inject_symmetry_lock(image_input, side) | |
| # Step B: EXTREME RESIZE (512px Max) | |
| # This is the only way to prevent Timeout on Free CPU | |
| w, h = processed_image.size | |
| MAX_SIZE = 512 | |
| scale = MAX_SIZE / max(w, h) | |
| new_w = int((w * scale) // 16 * 16) | |
| new_h = int((h * scale) // 16 * 16) | |
| print(f">>> Step 2: Resizing to {new_w}x{new_h} to prevent crash...") | |
| processed_image = processed_image.resize((new_w, new_h)) | |
| # Step C: Flux Inference | |
| print(f">>> Step 3: Running Flux (Prompt: {prompt})") | |
| generator = torch.Generator(DEVICE).manual_seed(int(seed)) | |
| try: | |
| result = pipe( | |
| prompt=prompt, | |
| image=processed_image, | |
| strength=strength, | |
| num_inference_steps=4, # Turbo Steps | |
| guidance_scale=0.0, | |
| generator=generator | |
| ).images[0] | |
| return result | |
| except Exception as e: | |
| print(f"ERROR: {e}") | |
| return None | |
| # 4. UI | |
| css = """ | |
| #col-container { max-width: 900px; margin: 0 auto; background-color: #f0f2f6; padding: 20px; border-radius: 10px; } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown("# ⚡ Flux Face Symmetry (Safe Mode)") | |
| gr.Markdown("Optimized for CPU Stability (Max 512px).") | |
| with gr.Row(): | |
| with gr.Column(): | |
| img_in = gr.Image(label="Upload Face", type="pil") | |
| side = gr.Radio(["Left", "Right"], label="Best Side", value="Left") | |
| prompt = gr.Text(label="Prompt", value="perfect symmetry, smooth skin") | |
| strength = gr.Slider(0.15, 0.45, value=0.25, step=0.01, label="Strength") | |
| seed = gr.Number(label="Seed", value=123) | |
| btn = gr.Button("Generate", variant="primary") | |
| with gr.Column(): | |
| img_out = gr.Image(label="Result") | |
| btn.click(process_image, inputs=[prompt, img_in, side, strength, seed], outputs=[img_out]) | |
| # Enable Queue to handle timeouts better | |
| demo.queue(max_size=2).launch() |