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()