app.py

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)