app.py

import os
# ZeroGPU runs on MIG slices; torch's expandable-segments allocator makes NVML
# calls that fail on MIG and surface as "NVML_SUCCESS == r INTERNAL ASSERT
# FAILED (CUDACachingAllocator)". Must be set before torch is imported.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False")

import random

import gradio as gr
import spaces
import torch
from diffusers import Krea2Pipeline
from transformers import AutoConfig, AutoModel, AutoTokenizer, Qwen2Tokenizer

# ----------------------------------------------------------------------------
# Krea-R-Turbo = Krea 2 Turbo + Rebels style LoRAs (strengths via Space secrets)
# LoRAs are fused at startup, so inference cost is identical to the base model.
# ----------------------------------------------------------------------------
BASE = "krea/Krea-2-Turbo"

# Krea's repo ships only tokenizer.json (fast format), but the pipeline's
# model_index declares the slow Qwen2Tokenizer, and diffusers' class check
# rejects a fast instance. So: load fast (with extra_special_tokens={} to
# dodge the list-vs-dict bug in transformers 4.x), export it — which writes
# the vocab.json/merges.txt the slow class needs — then load the slow class
# from the export. Verified: slow/fast produce identical token ids.
_fast = AutoTokenizer.from_pretrained(BASE, subfolder="tokenizer",
                                      use_fast=True, extra_special_tokens={})
_fast.save_pretrained("/tmp/krea_tokenizer")
tokenizer = Qwen2Tokenizer.from_pretrained("/tmp/krea_tokenizer",
                                           extra_special_tokens={})

# Krea's text_encoder config has rope_scaling: null; transformers 4.57's
# Qwen3-VL calls .get() on it and crashes. Inject the dict with the exact
# values 4.57 defaults to anyway, so the math is unchanged.
te_cfg = AutoConfig.from_pretrained(BASE, subfolder="text_encoder")
_txt = getattr(te_cfg, "text_config", te_cfg)
if getattr(_txt, "rope_scaling", None) is None:
    _txt.rope_scaling = {"rope_type": "default", "mrope_section": [24, 20, 20]}
text_encoder = AutoModel.from_pretrained(BASE, subfolder="text_encoder",
                                         config=te_cfg, torch_dtype=torch.bfloat16)

pipe = Krea2Pipeline.from_pretrained(BASE, tokenizer=tokenizer,
                                     text_encoder=text_encoder,
                                     torch_dtype=torch.bfloat16)

pipe.load_lora_weights("realrebelai/RebelReal_LoRA_Collection",
                       weight_name="RebelReal_(Krea-2).safetensors",
                       adapter_name="rebelreal")
pipe.load_lora_weights("realrebelai/RebelMidjourney_LoRA_Collection",
                       weight_name="RebelMidjourney_(Krea-2).safetensors",
                       adapter_name="rebelmj")
pipe.set_adapters(["rebelreal", "rebelmj"], adapter_weights=[0.31, 0.13])
# Bake the LoRAs into the weights (same math as an offline merge), then drop
# the adapter machinery so sampling has zero LoRA overhead.
pipe.fuse_lora(adapter_names=["rebelreal", "rebelmj"])
pipe.unload_lora_weights()

# ZeroGPU slices reload the packed weights on every call, so all-resident
# .to("cuda") needs weights+activations to fit at once (~35 GB + acts) and
# OOMs on smaller slices. Model offload keeps only the active component on
# GPU (text encoder -> transformer -> VAE), peak ~26 GB.
pipe.enable_model_cpu_offload()
# trim peak VRAM further: tiled VAE decode costs ~nothing at these resolutions
if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"):
    pipe.vae.enable_tiling()

# torch 2.11's fused SDPA kernels reject GQA (Krea2 has 48 query / 12 KV heads),
# so every fused path declines and attention falls to the math kernel (OOM) or
# nothing. Expand KV heads to match Q before dispatch: repeat_interleave is
# numerically identical to enable_gqa (verified bit-exact), costs ~120 MB.
import diffusers.models.transformers.transformer_krea2 as _tk

_orig_dispatch = _tk.dispatch_attention_fn

def _gqa_expanded_dispatch(query, key, value, *args, **kwargs):
    # layout here is [B, seq, heads, dim]; heads on dim 2
    hq, hkv = query.shape[2], key.shape[2]
    if hq != hkv:
        n = hq // hkv
        key = key.repeat_interleave(n, dim=2)
        value = value.repeat_interleave(n, dim=2)
        kwargs["enable_gqa"] = False
    return _orig_dispatch(query, key, value, *args, **kwargs)

_tk.dispatch_attention_fn = _gqa_expanded_dispatch

# Force the memory-efficient kernel: O(seq) memory, supports masks.
for _backend in ("_native_efficient", "_native_cudnn"):
    try:
        pipe.transformer.set_attention_backend(_backend)
        print(f"[attn] using backend: {_backend}", flush=True)
        break
    except Exception as _e:
        print(f"[attn] {_backend} unavailable: {_e}", flush=True)

MAX_SEED = 2**31 - 1


@spaces.GPU(duration=110)
def generate(prompt, width, height, steps, seed, randomize_seed,
             progress=gr.Progress(track_tqdm=True)):
    free_b, total_b = torch.cuda.mem_get_info()
    print(f"[vram] device={torch.cuda.get_device_name(0)} "
          f"total={total_b/1e9:.1f}GB free={free_b/1e9:.1f}GB", flush=True)
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator("cuda").manual_seed(int(seed))
    image = pipe(
        prompt=prompt,
        width=int(width),
        height=int(height),
        num_inference_steps=int(steps),
        guidance_scale=0.0,          # Turbo is CFG-free
        generator=generator,
    ).images[0]
    return image, seed


with gr.Blocks(title="Krea-R-Turbo") as demo:
    gr.Markdown(
        """
        # Krea-R-Turbo
        A custom merge of Krea-2-Turbo and 2 of Rebels style LoRAs at specific
        strength values. Displays a heavy focus on photorealistic portraits to
        achieve high grade aesthetics while also retaining Krea-2s sharp detail!

        GGUF quants for local ComfyUI (8 GB VRAM+):
        [realrebelai/Krea-R-Turbo](https://huggingface.co/realrebelai/Krea-R-Turbo)
        """
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=3,
                                placeholder="a cinematic photo of ...")
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Settings", open=False):
                with gr.Row():
                    width = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                    height = gr.Slider(512, 1536, value=1024, step=64, label="Height")
                steps = gr.Slider(4, 12, value=8, step=1, label="Steps")
                with gr.Row():
                    seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                    randomize_seed = gr.Checkbox(value=True, label="Random seed")
        with gr.Column():
            out = gr.Image(label="Result", format="png")
            used_seed = gr.Number(label="Seed used", interactive=False)

    run.click(generate,
              inputs=[prompt, width, height, steps, seed, randomize_seed],
              outputs=[out, used_seed])
    prompt.submit(generate,
                  inputs=[prompt, width, height, steps, seed, randomize_seed],
                  outputs=[out, used_seed])

demo.launch()