Update app.py

app.py

@@ -1 +1,71 @@
-import os; exec(os.getenv('EXEC'))
+import torch
+import gradio as gr
+import numpy as np
+from diffusers import FluxKontextPipeline
+from gfpgan import GFPGANer
+from PIL import Image
+
+# Device setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load FLUX in-context editing pipeline
+pipe = FluxKontextPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Kontext-dev",
+    torch_dtype=torch.bfloat16 if device.type=="cuda" else torch.float32
+).to(device)
+
+# Load face enhancement model
+gfpgan = GFPGANer(
+    model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.3.pth",
+    upscale=1,
+    arch="clean",
+    channel_multiplier=2,
+    bg_upsampler=None,
+    device=device.type
+)
+
+def enhance_face(input_img: Image.Image) -> Image.Image:
+    img_np = np.array(input_img)
+    _, _, output = gfpgan.enhance(img_np, has_aligned=False, only_center_face=False, paste_back=True)
+    return Image.fromarray(output)
+
+def infer(input_image, prompt, beautify, seed, randomize, steps, guidance_scale):
+    # Set random seed
+    generator = torch.Generator(device=device).manual_seed(seed if not randomize else torch.randint(0, 2**32-1, ()).item())
+
+    # In-context editing
+    out = pipe(
+        image=input_image.convert("RGB"),
+        prompt=prompt,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    ).images[0]
+
+    # Apply face enhancement if selected
+    if beautify:
+        out = enhance_face(out)
+
+    return out
+
+# UI setup
+with gr.Blocks() as demo:
+    gr.Markdown("# FLUX Kontekt Editor + Beautify")
+    with gr.Row():
+        input_image = gr.Image(label="Upload Image", type="pil")
+        result = gr.Image(label="Edited Output")
+    prompt = gr.Textbox(label="Edit Prompt", placeholder="e.g., 'change background to beach'")
+    beautify = gr.Checkbox(label="Beautify Face", value=True)
+    seed = gr.Slider(0, 2**32-1, value=0, step=1, label="Seed")
+    randomize = gr.Checkbox(label="Randomize Seed", value=True)
+    steps = gr.Slider(1, 30, value=28, label="Steps")
+    guidance = gr.Slider(1.0, 10.0, value=2.5, step=0.1, label="Guidance Scale")
+    run = gr.Button("Run")
+
+    run.click(
+        fn=infer,
+        inputs=[input_image, prompt, beautify, seed, randomize, steps, guidance],
+        outputs=result
+    )
+
+demo.launch()