import os import random import spaces import torch import gradio as gr from huggingface_hub import hf_hub_download from safetensors.torch import load_file as load_safetensors from diffusers import Krea2Pipeline TURBO_REPO = "unsloth/Krea-2-Turbo" HD_VAE_REPO = "wikeeyang/Krea2-Turbo-HD-V1" MAX_SEED = 2**31 - 1 def _remap_vae_keys(state_dict): """Remap ComfyUI-format VAE state dict keys to diffusers AutoencoderKLQwenImage keys. The ComfyUI checkpoint uses a flat sequential naming convention (residual.0, residual.2, etc.) while diffusers uses semantic names (norm1, conv1, norm2, conv2). This function translates between the two. """ # Mapping from ComfyUI decoder upsamples index to diffusers up_blocks path. # Derived from the AutoencoderKLQwenImage architecture with dim_mult=[1,2,4,4] # and num_res_blocks=2: each level has 2 resnets, with upsamplers between levels. up_map = { 0: "up_blocks.0.resnets.0", 1: "up_blocks.0.resnets.1", 2: "up_blocks.0.resnets.2", 3: "up_blocks.0.upsamplers.0", 4: "up_blocks.1.resnets.0", 5: "up_blocks.1.resnets.1", 6: "up_blocks.1.resnets.2", 7: "up_blocks.1.upsamplers.0", 8: "up_blocks.2.resnets.0", 9: "up_blocks.2.resnets.1", 10: "up_blocks.2.resnets.2", 11: "up_blocks.2.upsamplers.0", 12: "up_blocks.3.resnets.0", 13: "up_blocks.3.resnets.1", 14: "up_blocks.3.resnets.2", } def _fix_resnet(rest): """Map sequential residual indices to semantic names.""" rest = rest.replace("residual.0.", "norm1.") rest = rest.replace("residual.2.", "conv1.") rest = rest.replace("residual.3.", "norm2.") rest = rest.replace("residual.6.", "conv2.") rest = rest.replace("shortcut.", "conv_shortcut.") return rest def _fix_middle(key, side): """Map encoder/decoder middle.X to mid_block structure.""" parts = key.split(".") idx = int(parts[2]) rest = ".".join(parts[3:]) if idx == 0: # first resnet return f"{side}.mid_block.resnets.0.{_fix_resnet(rest)}" elif idx == 1: # attention return f"{side}.mid_block.attentions.0.{rest}" elif idx == 2: # second resnet return f"{side}.mid_block.resnets.1.{_fix_resnet(rest)}" new_state = {} for key, val in state_dict.items(): new_key = key # Top-level convs: quant_conv and post_quant_conv if key.startswith("conv1."): new_key = "quant_conv." + key[len("conv1."):] elif key.startswith("conv2."): new_key = "post_quant_conv." + key[len("conv2."):] # Encoder mapping elif key.startswith("encoder.conv1."): new_key = "encoder.conv_in." + key[len("encoder.conv1."):] elif key.startswith("encoder.head.0."): new_key = "encoder.norm_out." + key[len("encoder.head.0."):] elif key.startswith("encoder.head.2."): new_key = "encoder.conv_out." + key[len("encoder.head.2."):] elif key.startswith("encoder.downsamples."): parts = key.split(".") idx = int(parts[2]) rest = ".".join(parts[3:]) rest = _fix_resnet(rest) new_key = f"encoder.down_blocks.{idx}.{rest}" elif key.startswith("encoder.middle."): new_key = _fix_middle(key, "encoder") # Decoder mapping elif key.startswith("decoder.conv1."): new_key = "decoder.conv_in." + key[len("decoder.conv1."):] elif key.startswith("decoder.head.0."): new_key = "decoder.norm_out." + key[len("decoder.head.0."):] elif key.startswith("decoder.head.2."): new_key = "decoder.conv_out." + key[len("decoder.head.2."):] elif key.startswith("decoder.upsamples."): parts = key.split(".") idx = int(parts[2]) rest = ".".join(parts[3:]) rest = _fix_resnet(rest) new_key = f"decoder.{up_map[idx]}.{rest}" elif key.startswith("decoder.middle."): new_key = _fix_middle(key, "decoder") new_state[new_key] = val return new_state # Load the Krea-2-Turbo pipeline (full diffusers format) then swap in the # HD-optimized VAE from wikeeyang/Krea2-Turbo-HD-V1 for enhanced detail, # clarity, and contrast. The HD VAE is a ComfyUI-format single safetensors # checkpoint; we remap its keys to the diffusers AutoencoderKLQwenImage layout. pipe = Krea2Pipeline.from_pretrained(TURBO_REPO, torch_dtype=torch.bfloat16) # Load the HD VAE weights, remap keys from ComfyUI format, and replace the VAE hd_vae_path = hf_hub_download(HD_VAE_REPO, "Krea2-HD-vae.safetensors") hd_vae_state = load_safetensors(hd_vae_path) remapped = _remap_vae_keys(hd_vae_state) pipe.vae.load_state_dict(remapped, strict=True) print("HD VAE weights loaded and remapped successfully.") pipe.to("cuda") PROMPT_TIPS = """\ Krea 2 Turbo HD is tuned for natural language prompts. Describe the image as you would to a person. - Write in full sentences or rich phrases. Longer, more specific prompts give the best results. - Name the things that matter: subject, setting, lighting, color, framing, medium, and mood. - To render text in the image, wrap the words in quotes, e.g. a neon sign that reads "open late". - The HD-optimized VAE improves detail rendering, clarity, and contrast over the original. """ EXAMPLE_PROMPTS = [ [ "A mesmerizing painting that captures the essence of a dreamy, ethereal punk woman " "with delicate features. Her short hair frames her introspective face, reflecting the " "depth of emotion through her soft gaze. The harmonious background, a mix of soft, " "muted colors, creates a warm and intimate atmosphere. The delicate brushstrokes and " "vibrant color palette demonstrate the minimalist sophistication, subtle realism, " "and emotional depth of the artist., Mysterious, vibrant, painting, fashion." ], [ "This close-up shot captures an Eastern Gray Squirrel perched on a tree branch, " "its inquisitive gaze directed straight at the viewer. The squirrel's fur is a blend " "of gray and brown, with hints of reddish-brown on its bushy tail and flanks. Its " "underparts are a lighter, almost white color, providing a subtle contrast. The fur " "appears thick and soft, perfectly adapted for the colder climate it inhabits. The " "squirrel's eyes are large and dark, giving it an alert and curious expression." ], [ "Create a photorealistic concept poster featuring an elegant timepiece. The watch " "should be the focal point, displayed in a close-up view. The watch face is a " "high-end, luxury design, featuring intricate gold hands and markers that are " "sharply detailed. The reflection on the watch face should depict a gentle, golden " "sunset over a navy blue ocean. The ocean's surface should have subtle ripples, " "catching the last light of the day." ], [ "Super detailed 32K A majestic and ferocious tiger captured mid-leap, its powerful " "body cutting through a cascade of splashing water, embodying raw energy and primal " "strength. The tiger's face is at the center of the composition, its intense golden " "eyes fixed forward, radiating focus and dominance. The intricate black stripes on " "its fiery orange fur form mesmerizing patterns that emphasize its wild beauty and " "feral grace. Each strand of fur is meticulously detailed, shimmering with droplets " "of water that catch the light." ], ] PLACEHOLDER = ( "Describe your image in natural language. e.g. a russet harvest mouse clinging to a " 'branch, macro photograph, shallow depth of field, creamy green bokeh, soft natural light. ' 'Wrap words in "quotes" to render them as text.' ) CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ def _duration(prompt, width, height, steps, seed, randomize): """Estimate GPU time based on steps and pixel area.""" megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024)) return int(int(steps) * 2 * megapixels + 25) @spaces.GPU(duration=_duration, size="xlarge") def generate( prompt: str, width: int = 1024, height: int = 1024, steps: int = 8, seed: int = 42, randomize: bool = True, ): """Generate an image from a text prompt using Krea-2-Turbo with HD VAE. Args: prompt: Natural language description of the image to generate. width: Output image width in pixels. height: Output image height in pixels. steps: Number of inference steps (8 for Turbo). seed: RNG seed for reproducibility. randomize: If True, pick a random seed. """ if not prompt or not prompt.strip(): raise gr.Error("Enter a prompt to generate an image.") if randomize: seed = random.randint(0, MAX_SEED) seed = int(seed) generator = torch.Generator("cuda").manual_seed(seed) try: image = pipe( prompt=prompt, height=int(height), width=int(width), num_inference_steps=int(steps), guidance_scale=0.0, generator=generator, ).images[0] except RuntimeError as exc: torch.cuda.empty_cache() raise gr.Error( f"Generation failed at {int(width)}x{int(height)}. This is usually GPU " "out of memory. Try 1024x1024 or a smaller size." ) from exc return image, seed with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# Krea 2 Turbo HD\n" "Text-to-image generation with the HD-optimized VAE from " "[wikeeyang/Krea2-Turbo-HD-V1](https://huggingface.co/wikeeyang/Krea2-Turbo-HD-V1), " "built on [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo). " "The fine-tuned VAE enhances detail rendering, clarity, and contrast." ) with gr.Row(): prompt = gr.Textbox( show_label=False, placeholder=PLACEHOLDER, container=False, scale=4, lines=3, autofocus=True, ) run = gr.Button("Generate", variant="primary", scale=1) output = gr.Image(label="Result", format="png") with gr.Accordion("Advanced settings", open=False): with gr.Row(): width = gr.Slider(512, 2048, value=1024, step=16, label="Width") height = gr.Slider(512, 2048, value=1024, step=16, label="Height") steps = gr.Slider(1, 50, value=8, step=1, label="Steps") with gr.Row(): seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed") randomize = gr.Checkbox(value=True, label="Randomize seed") with gr.Accordion("Prompting tips", open=False): gr.Markdown(PROMPT_TIPS) gr.Examples( examples=EXAMPLE_PROMPTS, inputs=[prompt], outputs=[output, seed], fn=generate, cache_examples=True, cache_mode="lazy", label="Example prompts", ) run.click( generate, inputs=[prompt, width, height, steps, seed, randomize], outputs=[output, seed], api_name="generate", ) prompt.submit( generate, inputs=[prompt, width, height, steps, seed, randomize], outputs=[output, seed], ) demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)