import os import io import gc import uuid import json import base64 import random import subprocess from pathlib import Path from typing import List, Optional import spaces import numpy as np import torch from PIL import Image from gradio import Server from fastapi import Request, UploadFile, File, Form from fastapi.responses import HTMLResponse, JSONResponse, FileResponse HF_TOKEN = os.environ.get("HF_TOKEN") app = Server() BASE_DIR = Path(__file__).resolve().parent STATIC_DIR = BASE_DIR / "static" OUTPUT_DIR = BASE_DIR / "outputs" EXAMPLES_DIR = BASE_DIR / "examples" STATIC_DIR.mkdir(exist_ok=True) OUTPUT_DIR.mkdir(exist_ok=True) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 ADAPTER = { "title": "Klein-Consistency", "adapter_name": "klein-consistency", "repo": "dx8152/Flux2-Klein-9B-Consistency", "weights": "Klein-consistency.safetensors", } DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.bfloat16 if torch.cuda.is_available(): print("current device:", torch.cuda.current_device()) print("device name:", torch.cuda.get_device_name(torch.cuda.current_device())) DEVICE_LABEL = torch.cuda.get_device_name(torch.cuda.current_device()).lower() else: DEVICE_LABEL = str(DEVICE).lower() print("CUDA_VISIBLE_DEVICES =", os.environ.get("CUDA_VISIBLE_DEVICES")) print("torch.__version__ =", torch.__version__) print("Using device:", DEVICE) def apply_patch(): import diffusers site_packages = os.path.dirname(diffusers.__file__) patch_file = os.path.join(os.path.dirname(__file__), "flux2_klein_kv.patch") if os.path.exists(patch_file): result = subprocess.run( ["patch", "-p2", "--forward", "--batch"], cwd=os.path.dirname(site_packages), stdin=open(patch_file), capture_output=True, text=True, ) if result.returncode == 0: print("Patch applied successfully") else: print(f"Patch output: {result.stdout}\n{result.stderr}") apply_patch() from diffusers.pipelines.flux2.pipeline_flux2_klein_kv import Flux2KleinKVPipeline print("Loading FLUX.2 Klein 9B KV model...") pipe = Flux2KleinKVPipeline.from_pretrained( "black-forest-labs/FLUX.2-klein-9b-kv", torch_dtype=dtype, token=HF_TOKEN, ).to(DEVICE) print("Base KV Model loaded successfully.") print(f"Loading adapter: {ADAPTER['title']}") pipe.load_lora_weights( ADAPTER["repo"], weight_name=ADAPTER["weights"], adapter_name=ADAPTER["adapter_name"], ) pipe.set_adapters([ADAPTER["adapter_name"]], adapter_weights=[1.0]) print(f"Adapter loaded successfully: {ADAPTER['adapter_name']}") def image_to_base64(img: Image.Image) -> str: buf = io.BytesIO() img.save(buf, format="PNG") return base64.b64encode(buf.getvalue()).decode("utf-8") def save_image(img: Image.Image, prefix: str = "output") -> str: filename = f"{prefix}_{uuid.uuid4().hex}.png" path = OUTPUT_DIR / filename img.save(path, format="PNG") return filename def update_dimensions_on_upload(image): if image is None: return 1024, 1024 try: if isinstance(image, list) and len(image) > 0: first = image[0] else: first = image if isinstance(first, (tuple, list)): path_or_img = first[0] else: path_or_img = first if isinstance(path_or_img, str): img = Image.open(path_or_img).convert("RGB") elif isinstance(path_or_img, Image.Image): img = path_or_img.convert("RGB") else: img = Image.open(path_or_img.name).convert("RGB") original_width, original_height = img.size if original_width > original_height: new_width = 1024 aspect_ratio = original_height / original_width new_height = int(new_width * aspect_ratio) else: new_height = 1024 aspect_ratio = original_width / original_height new_width = int(new_height * aspect_ratio) new_width = (new_width // 8) * 8 new_height = (new_height // 8) * 8 new_width = max(256, min(1024, new_width)) new_height = max(256, min(1024, new_height)) return new_width, new_height except Exception: return 1024, 1024 def process_gallery_images(images): if not images: return [] pil_images = [] for item in images: try: if isinstance(item, (tuple, list)): path_or_img = item[0] else: path_or_img = item if isinstance(path_or_img, str): pil_images.append(Image.open(path_or_img).convert("RGB")) elif isinstance(path_or_img, Image.Image): pil_images.append(path_or_img.convert("RGB")) else: pil_images.append(Image.open(path_or_img.name).convert("RGB")) except Exception as e: print(f"Skipping invalid image item: {e}") continue return pil_images @spaces.GPU def infer( images, prompt, seed, randomize_seed, width, height, steps, ): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if not prompt or not str(prompt).strip(): raise ValueError("Please enter a prompt.") if isinstance(seed, str): seed = int(seed) if isinstance(randomize_seed, str): randomize_seed = randomize_seed.lower() == "true" if isinstance(width, str): width = int(width) if isinstance(height, str): height = int(height) if isinstance(steps, str): steps = int(steps) if randomize_seed: seed = random.randint(0, MAX_SEED) pil_images = process_gallery_images(images) if images else [] if pil_images: width, height = update_dimensions_on_upload(pil_images[0]) image_input = [ img.resize((width, height), Image.LANCZOS).convert("RGB") for img in pil_images ] else: image_input = None width = max(256, min(MAX_IMAGE_SIZE, (int(width) // 8) * 8)) height = max(256, min(MAX_IMAGE_SIZE, (int(height) // 8) * 8)) try: generator = torch.Generator(device=DEVICE).manual_seed(seed) pipe_kwargs = { "prompt": prompt, "width": width, "height": height, "num_inference_steps": steps, "generator": generator, } if image_input is not None: pipe_kwargs["image"] = image_input result_image = pipe(**pipe_kwargs).images[0] return result_image, seed except Exception as e: raise RuntimeError(f"Inference failed: {e}") finally: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def get_example_items(): example_prompts = { "1.jpg": "Change the weather to stormy.", "2.jpg": "Transform the scene into a snowy winter day while preserving the original subject identity, framing, and composition.", "3.jpg": "Relight the image with soft golden sunset lighting while keeping all structures and subject details consistent.", "4.jpg": "Make the texture high-resolution.", } items = [] if EXAMPLES_DIR.exists(): for name in sorted(os.listdir(EXAMPLES_DIR)): if name.lower().endswith((".png", ".jpg", ".jpeg", ".webp")): items.append( { "file": name, "url": f"/example-file/{name}", "prompt": example_prompts.get(name, "Edit this image while preserving composition."), } ) return items @app.api(name="hello") def hello(name: str) -> str: return f"Hello, {name}!" @app.get("/example-file/{filename}") async def example_file(filename: str): path = EXAMPLES_DIR / filename if not path.exists(): return JSONResponse({"error": "Example not found"}, status_code=404) return FileResponse(path) @app.get("/download/{filename}") async def download_file(filename: str): path = OUTPUT_DIR / filename if not path.exists(): return JSONResponse({"error": "File not found"}, status_code=404) return FileResponse(path, filename=filename, media_type="image/png") @app.post("/api/edit") async def edit_image( prompt: str = Form(...), seed: str = Form("0"), randomize_seed: str = Form("true"), width: str = Form("1024"), height: str = Form("1024"), steps: str = Form("4"), images: Optional[List[UploadFile]] = File(None), ): temp_paths = [] try: image_paths = [] if images: for upload in images: suffix = Path(upload.filename).suffix or ".png" temp_name = f"upload_{uuid.uuid4().hex}{suffix}" temp_path = OUTPUT_DIR / temp_name content = await upload.read() with open(temp_path, "wb") as f: f.write(content) temp_paths.append(str(temp_path)) image_paths.append(str(temp_path)) result_image, used_seed = infer( images=image_paths, prompt=prompt, seed=seed, randomize_seed=randomize_seed, width=width, height=height, steps=steps, ) output_filename = save_image(result_image, prefix="kv_edit") return JSONResponse( { "success": True, "seed": used_seed, "image_url": f"/download/{output_filename}", "download_url": f"/download/{output_filename}", "image_base64": image_to_base64(result_image), "device": DEVICE_LABEL, } ) except Exception as e: return JSONResponse( {"success": False, "error": str(e)}, status_code=500, ) finally: for p in temp_paths: try: if os.path.exists(p): os.remove(p) except Exception: pass @app.get("/", response_class=HTMLResponse) async def homepage(request: Request): examples = get_example_items() examples_json = json.dumps(examples) return f""" KV-Edit-Consistency
4-Step Fast Inference KV Image Editing Playground
black-forest-labs / flux.2-klein-9b-kv / edit

KV-Edit-Consistency

Inference
image-to-image
fast-edit

Input

Form
Images
The first uploaded image is used to auto-fit width and height while preserving aspect ratio.

Result

Idle
No output yet
Your edited image will appear here
Generated output
Processing image
seed
-
device name
{DEVICE_LABEL}
Examples
""" app.launch()