import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import random import re import spaces import torch import gradio as gr from huggingface_hub import hf_hub_download from safetensors.torch import load_file from diffusers import Krea2Pipeline # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- # The Enhancer LoRA (vrgamedevgirl84/Krea2_Enhancer) was trained on Krea 2 RAW. # Krea recommends training on RAW and applying on Turbo for fast inference, and # the LoRA applies cleanly to Turbo. Turbo (8-step distilled) makes this a # snappy ZeroGPU demo, matching the proven krea/krea-lora-the-explorer pattern. DTYPE = torch.bfloat16 BASE_MODEL = "krea/Krea-2-Turbo" LORA_REPO = "vrgamedevgirl84/Krea2_Enhancer" LORA_FILE = "krea2_Enhancer.safetensors" ADAPTER_NAME = "enhancer" MAX_SEED = 2**31 - 1 TOKEN = os.environ.get("HF_TOKEN") # --------------------------------------------------------------------------- # Kohya (musubi-tuner) -> diffusers key conversion # --------------------------------------------------------------------------- # The Enhancer LoRA ships in kohya/musubi format (ss_network_module: # networks.lora_krea2), with keys like `lora_unet_blocks_0_attn_wq.lora_down`. # diffusers' Krea2 loader expects `transformer.transformer_blocks.0.attn.to_q. # lora_A`. Every one of the 264 modules maps 1:1 with identical shapes; alpha # equals rank (32), so the effective scale is 1.0 and no rescale is needed. _SUB = { "attn_wq": "attn.to_q", "attn_wk": "attn.to_k", "attn_wv": "attn.to_v", "attn_wo": "attn.to_out.0", "attn_gate": "attn.to_gate", "mlp_down": "ff.down", "mlp_gate": "ff.gate", "mlp_up": "ff.up", } _SPECIAL = { "lora_unet_first": "transformer.img_in", "lora_unet_last_linear": "transformer.final_layer.linear", "lora_unet_tmlp_0": "transformer.time_embed.linear_1", "lora_unet_tmlp_2": "transformer.time_embed.linear_2", "lora_unet_tproj_1": "transformer.time_mod_proj", "lora_unet_txtmlp_1": "transformer.txt_in.linear_1", "lora_unet_txtmlp_3": "transformer.txt_in.linear_2", "lora_unet_txtfusion_projector": "transformer.text_fusion.projector", } def _kohya_module_to_diffusers(mod: str): if mod in _SPECIAL: return _SPECIAL[mod] m = re.match(r"lora_unet_blocks_(\d+)_(.+)", mod) if m: return f"transformer.transformer_blocks.{m.group(1)}.{_SUB[m.group(2)]}" m = re.match(r"lora_unet_txtfusion_layerwise_blocks_(\d+)_(.+)", mod) if m: return f"transformer.text_fusion.layerwise_blocks.{m.group(1)}.{_SUB[m.group(2)]}" m = re.match(r"lora_unet_txtfusion_refiner_blocks_(\d+)_(.+)", mod) if m: return f"transformer.text_fusion.refiner_blocks.{m.group(1)}.{_SUB[m.group(2)]}" return None def _convert_kohya_state_dict(sd): """Convert a kohya Krea2 LoRA state dict to diffusers lora_A/lora_B keys.""" # collect per-module alpha to fold into scaling (alpha/rank); here alpha==rank. alphas = {k[: -len(".alpha")]: v for k, v in sd.items() if k.endswith(".alpha")} out = {} skipped = [] for key, tensor in sd.items(): if key.endswith(".alpha"): continue if key.endswith(".lora_down.weight"): mod, suffix = key[: -len(".lora_down.weight")], "lora_A" elif key.endswith(".lora_up.weight"): mod, suffix = key[: -len(".lora_up.weight")], "lora_B" else: skipped.append(key) continue target = _kohya_module_to_diffusers(mod) if target is None: skipped.append(key) continue w = tensor.to(DTYPE) # fold alpha/rank into lora_B (diffusers applies scale * B @ A) if suffix == "lora_B" and mod in alphas: rank = tensor.shape[1] alpha = float(alphas[mod].float().item()) scale = alpha / rank if abs(scale - 1.0) > 1e-6: w = w * scale out[f"{target}.{suffix}.weight"] = w if skipped: print(f"[lora] skipped {len(skipped)} keys (e.g. {skipped[:3]})") print(f"[lora] converted {len(out)} tensors to diffusers format") return out # --------------------------------------------------------------------------- # Load base pipeline + Enhancer LoRA at module scope (ZeroGPU packs at startup) # --------------------------------------------------------------------------- pipe = Krea2Pipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE) _lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=TOKEN) _kohya_sd = load_file(_lora_path) _diffusers_sd = _convert_kohya_state_dict(_kohya_sd) pipe.load_lora_weights(_diffusers_sd, adapter_name=ADAPTER_NAME) print("[lora] Enhancer LoRA loaded into transformer") pipe.to("cuda") # Track the scale currently baked into the adapter so we only re-set on change. _CURRENT_SCALE = {"value": None} def _set_scale(scale: float): if _CURRENT_SCALE["value"] != scale: pipe.set_adapters(ADAPTER_NAME, adapter_weights=float(scale)) _CURRENT_SCALE["value"] = scale # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=90, size="large") def generate( prompt: str, lora_scale: float = 1.0, steps: int = 8, guidance: float = 0.0, width: int = 1024, height: int = 1024, seed: int = 0, randomize_seed: bool = True, progress=gr.Progress(track_tqdm=True), ): """Generate an enhanced image from a text prompt using Krea 2 Turbo + the Krea2 Enhancer LoRA (sharper detail, cleaner textures, better lighting). Args: prompt: Text description of the image to generate. lora_scale: Strength of the Enhancer LoRA (0 = base model, 1 = full). steps: Number of denoising steps (Turbo works well at 8). guidance: Classifier-free guidance scale (Turbo uses 0). width: Image width in pixels. height: Image height in pixels. seed: RNG seed for reproducibility. randomize_seed: If true, ignore seed and pick a random one. """ if not prompt or not prompt.strip(): raise gr.Error("Please enter a prompt.") if randomize_seed: seed = random.randint(0, MAX_SEED) seed = int(seed) _set_scale(float(lora_scale)) generator = torch.Generator("cuda").manual_seed(seed) image = pipe( prompt=prompt.strip(), num_inference_steps=int(steps), guidance_scale=float(guidance), width=int(width), height=int(height), generator=generator, ).images[0] return image, seed # --------------------------------------------------------------------------- # UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } #result-image { min-height: 420px; } """ EXAMPLE_PROMPTS = [ ["A cinematic portrait of a woman in golden hour light, freckles, soft bokeh, ultra detailed skin"], ["A cozy wooden cabin in a snowy pine forest at dusk, warm glowing windows, volumetric light"], ["Close-up product shot of a luxury wristwatch on black marble, studio lighting, reflections"], ["An anime-style hero standing on a cliff overlooking a neon city, dramatic clouds, vivid colors"], ["A photorealistic bowl of ramen with steam, chopsticks, moody restaurant lighting, shallow depth of field"], ["A fantasy castle floating among clouds at sunrise, intricate architecture, epic scale"], ] with gr.Blocks(title="Krea 2 Enhancer") as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # 🎨 Krea 2 Enhancer Text-to-image with the **[Krea2 Enhancer LoRA](https://huggingface.co/vrgamedevgirl84/Krea2_Enhancer)** on top of **[Krea 2 Turbo](https://huggingface.co/krea/Krea-2-Turbo)** (8-step). The Enhancer adds cleaner detail, sharper textures, better lighting, and a more polished look across any style — while keeping your prompt intact. """ ) with gr.Row(): prompt = gr.Textbox( label="Prompt", lines=2, scale=4, placeholder="Describe the image you want to generate…", ) run_button = gr.Button("Generate", variant="primary", scale=1) result = gr.Image(label="Result", format="png", elem_id="result-image") with gr.Accordion("Advanced settings", open=False): lora_scale = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="Enhancer strength") steps = gr.Slider(1, 20, value=8, step=1, label="Steps") guidance = gr.Slider(0.0, 6.0, value=0.0, step=0.1, label="Guidance scale") with gr.Row(): width = gr.Slider(512, 1536, value=1024, step=16, label="Width") height = gr.Slider(512, 1536, value=1024, step=16, label="Height") with gr.Row(): seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") randomize_seed = gr.Checkbox(value=True, label="Randomize seed") gr.Examples( examples=EXAMPLE_PROMPTS, inputs=[prompt], outputs=[result, seed], fn=generate, cache_examples=True, cache_mode="lazy", ) inputs = [prompt, lora_scale, steps, guidance, width, height, seed, randomize_seed] gr.on( [run_button.click, prompt.submit], generate, inputs=inputs, outputs=[result, seed], api_name="generate", ) demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)