Spaces:
Running on Zero
Running on Zero
| import glob | |
| import os | |
| import random | |
| import subprocess | |
| import sys | |
| import spaces | |
| import torch | |
| import gradio as gr | |
| from huggingface_hub import login | |
| if os.environ.get("HF_TOKEN"): | |
| login(token=os.environ["HF_TOKEN"]) | |
| from diffusers import Krea2Pipeline | |
| DTYPE = torch.bfloat16 | |
| RAW_REPO = "krea/Krea-2-Raw" | |
| TURBO_REPO = "krea/Krea-2-Turbo" | |
| MAX_SEED = 2**31 - 1 | |
| # Both models are loaded at global scope. They share the architecture, so Turbo | |
| # reuses Raw's text encoder / tokenizer / VAE and only loads its own transformer. | |
| pipe_raw = Krea2Pipeline.from_pretrained(RAW_REPO, torch_dtype=DTYPE) | |
| pipe_turbo = Krea2Pipeline.from_pretrained( | |
| TURBO_REPO, | |
| text_encoder=pipe_raw.text_encoder, | |
| tokenizer=pipe_raw.tokenizer, | |
| vae=pipe_raw.vae, | |
| torch_dtype=DTYPE, | |
| ) | |
| pipe_raw.to("cuda") | |
| pipe_turbo.to("cuda") | |
| def _load_aoti(): | |
| # One compiled Krea2TransformerBlock (kernels only, weights stay live) serves | |
| # both pipelines since their transformer configs are identical. Mirrors | |
| # spaces.aoti_blocks_load, but downloads from the private artifact dataset | |
| # with an explicit write token instead of ambient model-repo auth. | |
| from huggingface_hub import hf_hub_download | |
| from spaces.zero.torch.aoti import LazyAOTIModel | |
| pt2 = hf_hub_download( | |
| repo_id="multimodalart/Krea-2-aoti", | |
| filename="Krea2TransformerBlock/package.pt2", | |
| repo_type="dataset", | |
| token=os.environ.get("HF_WRITE_TOKEN"), | |
| ) | |
| aoti_model = LazyAOTIModel(pt2) | |
| for pipe in (pipe_raw, pipe_turbo): | |
| for block in pipe.transformer.modules(): | |
| if block.__class__.__name__ == "Krea2TransformerBlock": | |
| spaces.aoti_patch(block, aoti_model) | |
| try: | |
| _load_aoti() | |
| print("AoTI blocks loaded.") | |
| except Exception as e: | |
| print(f"AoTI load skipped ({e}); running eager.") | |
| PIPES = {"Raw": pipe_raw, "Turbo": pipe_turbo} | |
| DEFAULTS = { | |
| "Raw": {"steps": 28, "guidance": 4.5}, | |
| "Turbo": {"steps": 8, "guidance": 0.0}, | |
| } | |
| # "Prompt strength" — amplifies and rebalances the positive conditioning to | |
| # counteract the trained safety filter's quality dilution. The tensor has shape | |
| # (B, seq, 12, 2560) — 12 Qwen3-VL encoder taps stacked on their own axis (dim | |
| # -2), not flattened into the feature dim. Per-layer weights boost the | |
| # middle/late taps (indices 7, 8, 10) that the filter suppresses most, while the | |
| # global strength scales the whole embedding. The defaults match the recommended | |
| # values from ComfyUI-Conditioning-Krea2Rebalance. | |
| PROMPT_STRENGTH_MIN = 0.5 | |
| PROMPT_STRENGTH_MAX = 4.0 | |
| PROMPT_STRENGTH_DEFAULT = 1.0 | |
| # Resolution presets. The model renders up to 2K, but the compiled transformer | |
| # block can exceed this Space's GPU memory above 1024, so 1024 is the default | |
| # and larger sizes are opt-in (see the OOM guard in generate). | |
| RESOLUTIONS = { | |
| "Square · 1024": (1024, 1024), | |
| "Portrait · 1024": (832, 1216), | |
| "Landscape · 1024": (1216, 832), | |
| "Square · 2K": (2048, 2048), | |
| } | |
| PROMPT_TIPS = """\ | |
| Krea 2 is tuned for natural language. Describe the image the way you would describe it to a person. | |
| - Write in full sentences or rich phrases. Longer, more specific prompts give the best results, but short prompts work too. | |
| - Name the things that matter: subject, setting, lighting, color, framing, medium, and mood. | |
| - To render text in the image, wrap the words in quotes, for example: a storefront window with a neon sign that reads "open late". | |
| - The model can render up to 2K, but very high resolutions may run out of GPU memory on this Space. 1024 is the reliable default. | |
| Want help writing longer prompts? An `expansion.txt` system prompt is provided in the [model repo](https://huggingface.co/krea/Krea-2-Turbo) for use with any LLM. | |
| """ | |
| # Drawn from the official Krea 2 prompt guide. These demonstrate the | |
| # detailed, natural-language style the model rewards. | |
| EXAMPLE_PROMPTS = [ | |
| ["immense rocket launch exhaust as seen from extremely close up"], | |
| [ | |
| "3D rendered matte black designer toy figure, stylized round anthropomorphic shape, " | |
| "backward black baseball cap, oversized gold-rimmed aviator sunglasses, white traditional " | |
| "line-art tattoos of tiger and bird on torso, black studded belt with gold buckle, smooth " | |
| "vinyl texture, studio lighting, solid vibrant blue background, high contrast minimal composition" | |
| ], | |
| [ | |
| "A tiny, russet-brown harvest mouse clings to a slender diagonal branch amid vibrant green " | |
| "lobed leaves and small round buds. The mouse has soft textured fur, glossy black eyes, a pink " | |
| "nose, fine whiskers, and delicate pink paws firmly gripping the wood. In this macro photograph, " | |
| "an extremely shallow depth of field sharply focuses on the animal's face. The deep green " | |
| "background dissolves into a smooth, creamy bokeh, illuminated by soft, diffused natural lighting " | |
| "that highlights the intricate details of the fur and foliage." | |
| ], | |
| [ | |
| "high-fashion editorial portrait of a young East Asian woman, short choppy platinum blonde bob " | |
| "with heavy bangs, looking over her bare shoulder to the right, lips playfully pursed, wearing a " | |
| "structured black top with an architectural protruding bust detail and thin straps, delicate gold " | |
| "hoop earrings, arm bent with hand resting on hip, warm skin tones, solid striking crimson red " | |
| "background, soft directional studio lighting, cinematic color palette, medium close-up shot" | |
| ], | |
| [ | |
| "A minimalist flat-color illustration of a person wading through expansive shallow ocean waves " | |
| "beneath a pale peach sky. The dark-skinned figure, wearing an orange swim cap, light blue top, and " | |
| "bright green shorts, steps carefully through knee-deep water. The ocean is rendered in muted mint " | |
| "green with delicate, thin black linework detailing the continuous ripples and gentle whitecaps. " | |
| "Soft pinkish-peach reflections echo the sky on the water's surface. The high-angle wide perspective " | |
| "emphasizes the vast negative space of the water, utilizing a clean ligne claire drawing aesthetic " | |
| "with a subtle paper texture." | |
| ], | |
| [ | |
| "A surreal retro-futuristic space scene features liquid chrome forming an abstract face merging " | |
| "with a glowing planetary horizon. The foreground is dominated by swirling, highly reflective " | |
| "metallic fluid that distorts into a stylized, melting facial profile with deep shadows and bright " | |
| "silver highlights. This undulating chrome form rests against the curved, atmospheric edge of a " | |
| "massive planet bathed in a soft electric blue and purple glow. Set against a deep black starfield, " | |
| "the artwork employs a vintage 1980s airbrush aesthetic with smooth gradients, ethereal lighting, " | |
| "and high-contrast metallic rendering." | |
| ], | |
| [ | |
| "Stylized digital painting of a menacing jester figure rendered with bold, expressive brushstrokes " | |
| "and a vibrant, almost psychedelic color palette against a pitch-black background. Dynamic low-angle " | |
| "perspective forces a dramatic, imposing composition as the character leans forward, one leg raised " | |
| "high. The jester wears a classic multi-pointed hat with bells, a ruffled collar, and striped tights " | |
| "in alternating shades of purple, blue, and chartreuse. The figure's face is a smooth, faceless, pale " | |
| "mauve mask with a single glowing white point of light at the center, and it grips a massive ornate " | |
| "sword with a glowing ethereal white blade. Theatrical lighting, dark fantasy concept-art aesthetic." | |
| ], | |
| [ | |
| "A close-up portrait of a young East Asian woman with straight black hair, loose strands sweeping " | |
| "across her fair skin, and an intense gaze. She wears a light grey collared shirt with a black tie. " | |
| "A vibrant bouquet of pink and orange lilies with lush green leaves sits in the blurred right " | |
| "foreground. The background is a solid, striking crimson red. Soft, directional studio lighting " | |
| "highlights her facial features, creating a high-contrast composition with a shallow depth of field." | |
| ], | |
| ] | |
| # Official sample renders for the prompts above, in the same order. Drop the | |
| # matching PNGs into assets/samples/ and the gallery shows them; clicking a | |
| # thumbnail loads its prompt. Filenames follow the Krea 2 prompt guide. If the | |
| # files are absent, the UI falls back to the text example prompts below. | |
| SAMPLE_DIR = os.path.join(os.path.dirname(__file__), "assets", "samples") | |
| SAMPLE_FILES = ["takeoff.png", "3d.png", "mouse.png", "red.png", "beach.png", "future.png", "jester.png", "flowers.png"] | |
| SAMPLE_LABELS = [ | |
| "Rocket exhaust", | |
| "Designer toy", | |
| "Harvest mouse", | |
| "Editorial portrait", | |
| "Ligne claire beach", | |
| "Liquid chrome", | |
| "Jester", | |
| "Crimson portrait", | |
| ] | |
| _gallery = [ | |
| (os.path.join(SAMPLE_DIR, fname), label, prompt[0]) | |
| for fname, label, prompt in zip(SAMPLE_FILES, SAMPLE_LABELS, EXAMPLE_PROMPTS) | |
| if os.path.exists(os.path.join(SAMPLE_DIR, fname)) | |
| ] | |
| GALLERY_ITEMS = [(path, label) for path, label, _ in _gallery] | |
| GALLERY_PROMPTS = [prompt for _, _, prompt in _gallery] | |
| 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.' | |
| ) | |
| def _duration(prompt, negative_prompt, model, steps, guidance, width, height, seed, randomize, | |
| prompt_strength=PROMPT_STRENGTH_DEFAULT, progress=None): | |
| # Scale the GPU reservation by step count and pixel area so larger renders | |
| # are not killed before they finish. | |
| megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024)) | |
| return int(int(steps) * 2 * megapixels + 25) | |
| def _encode_scaled(pipe, prompt, strength): | |
| """Encode `prompt` and apply per-layer rebalancing + global scaling. | |
| Krea 2 conditioning arrives as (B, seq, num_layers, dim) — the Qwen3-VL | |
| encoder taps are stacked along their own axis (dim -2), not flattened into | |
| the feature dim. We weight each tap individually to counteract the safety | |
| filter's quality dilution, then apply the global strength. | |
| Default per-layer weights from ComfyUI-Conditioning-Krea2Rebalance: boosting | |
| taps 7 (2.5x), 8 (5.0x), and 10 (4.0x) recovers detail the filter suppresses. | |
| """ | |
| prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( | |
| prompt=prompt, | |
| device="cuda", | |
| num_images_per_prompt=1, | |
| ) | |
| # Per-layer gains: taps 7, 8, 10 boosted to recover filter-suppressed detail. | |
| gains = torch.tensor( | |
| [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.5, 5.0, 1.1, 4.0, 1.0], | |
| dtype=torch.float32, device=prompt_embeds.device, | |
| ) | |
| orig_dtype = prompt_embeds.dtype | |
| t = prompt_embeds.float() | |
| n_layers = t.shape[-2] | |
| # Taps live on dim -2; scale each by its gain, broadcasting over the feature | |
| # dim. Skip the rebalance if the tap count ever diverges from the tuned | |
| # 12-entry table so a model/config change can't crash encoding. | |
| if gains.numel() == n_layers: | |
| t = t * gains.view(n_layers, 1) | |
| return (t * float(strength)).to(orig_dtype), prompt_embeds_mask | |
| def generate( | |
| prompt, | |
| negative_prompt="", | |
| model="Turbo", | |
| steps=8, | |
| guidance=0, | |
| width=1024, | |
| height=1024, | |
| seed=42, | |
| randomize=False, | |
| prompt_strength=PROMPT_STRENGTH_DEFAULT, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| 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) | |
| pipe = PIPES[model] | |
| strength = max(PROMPT_STRENGTH_MIN, min(PROMPT_STRENGTH_MAX, float(prompt_strength))) | |
| prompt_embeds, prompt_embeds_mask = _encode_scaled(pipe, prompt, strength) | |
| cond_kwargs = {"prompt_embeds": prompt_embeds, "prompt_embeds_mask": prompt_embeds_mask} | |
| try: | |
| image = pipe( | |
| negative_prompt=(negative_prompt or None) if guidance > 0 else None, | |
| height=int(height), | |
| width=int(width), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(guidance), | |
| generator=generator, | |
| **cond_kwargs, | |
| ).images[0] | |
| except RuntimeError as exc: | |
| # At high resolution the compiled transformer block can exhaust GPU | |
| # memory, which surfaces as a CUDA allocation / AOTI runtime error. | |
| # Recover the worker and tell the user how to fix it. | |
| torch.cuda.empty_cache() | |
| raise gr.Error( | |
| f"Generation failed at {int(width)}x{int(height)}. This is usually the GPU running " | |
| "out of memory at high resolution. Try 1024x1024 or a smaller size." | |
| ) from exc | |
| return image, seed | |
| def on_model_change(model): | |
| d = DEFAULTS[model] | |
| return ( | |
| gr.update(value=d["steps"]), | |
| gr.update(value=d["guidance"]), | |
| gr.update(interactive=d["guidance"] > 0), | |
| ) | |
| def on_resolution_change(label): | |
| w, h = RESOLUTIONS[label] | |
| return gr.update(value=w), gr.update(value=h) | |
| # Krea brand identity: neutral grayscale foundation with a single blue action | |
| # accent (krea.ai/press). Dark surfaces, mono utility type, accent reserved for | |
| # the primary action and focus states. | |
| KREA_ACCENT = "#2b5cff" | |
| theme = gr.themes.Base( | |
| primary_hue=gr.themes.colors.blue, | |
| neutral_hue=gr.themes.colors.neutral, | |
| font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], | |
| font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"], | |
| ).set( | |
| body_background_fill="#000000", | |
| body_background_fill_dark="#000000", | |
| body_text_color="#f5f5f5", | |
| background_fill_primary="#0d0d0d", | |
| background_fill_secondary="#0d0d0d", | |
| block_background_fill="#0d0d0d", | |
| block_border_color="#262626", | |
| block_border_width="1px", | |
| block_label_background_fill="#0d0d0d", | |
| block_label_text_color="#737373", | |
| block_title_text_color="#d4d4d5", | |
| border_color_primary="#262626", | |
| input_background_fill="#000000", | |
| input_border_color="#262626", | |
| input_border_color_focus=KREA_ACCENT, | |
| button_primary_background_fill=KREA_ACCENT, | |
| button_primary_background_fill_hover="#1f4fff", | |
| button_primary_text_color="#ffffff", | |
| button_primary_border_color=KREA_ACCENT, | |
| button_secondary_background_fill="#171717", | |
| button_secondary_background_fill_hover="#262626", | |
| button_secondary_text_color="#f5f5f5", | |
| button_secondary_border_color="#262626", | |
| slider_color=KREA_ACCENT, | |
| ) | |
| CSS = """ | |
| .gradio-container { background: #000 !important; } | |
| #page { max-width: 1120px; margin: 0 auto; padding: 4px 8px 32px; } | |
| #krea-header { | |
| padding: 32px 6px 22px; | |
| border-bottom: 1px solid #1a1a1a; | |
| margin-bottom: 22px; | |
| } | |
| #krea-header .eyebrow { | |
| font-family: 'JetBrains Mono', ui-monospace, monospace; | |
| font-size: 11px; | |
| letter-spacing: 0.24em; | |
| text-transform: uppercase; | |
| color: #737373; | |
| } | |
| #krea-header h1 { | |
| font-size: 42px; | |
| font-weight: 600; | |
| letter-spacing: -0.025em; | |
| line-height: 1.05; | |
| margin: 10px 0 6px; | |
| color: #fff; | |
| } | |
| #krea-header .subtitle { | |
| font-size: 15px; | |
| line-height: 1.5; | |
| color: #a3a3a3; | |
| margin: 0; | |
| max-width: 60ch; | |
| } | |
| #krea-header .meta { | |
| margin-top: 18px; | |
| display: flex; | |
| justify-content: space-between; | |
| align-items: center; | |
| flex-wrap: wrap; | |
| gap: 12px; | |
| } | |
| #krea-header .badges { display: flex; gap: 8px; } | |
| #krea-header .badge { | |
| font-family: 'JetBrains Mono', ui-monospace, monospace; | |
| font-size: 10px; | |
| letter-spacing: 0.12em; | |
| text-transform: uppercase; | |
| color: #d4d4d5; | |
| border: 1px solid #262626; | |
| border-radius: 999px; | |
| padding: 4px 10px; | |
| } | |
| #krea-header .links { display: flex; gap: 16px; } | |
| #krea-header .links a { | |
| font-family: 'JetBrains Mono', ui-monospace, monospace; | |
| font-size: 11px; | |
| letter-spacing: 0.08em; | |
| text-transform: uppercase; | |
| color: #737373; | |
| text-decoration: none; | |
| transition: color 0.15s ease; | |
| } | |
| #krea-header .links a:hover { color: #f5f5f5; } | |
| #generate-btn { font-weight: 600; letter-spacing: 0.01em; } | |
| #result-image { min-height: 420px; border-radius: 10px; overflow: hidden; } | |
| /* Inline code chips legible on the dark theme (e.g. the expansion.txt mention in tips). */ | |
| .gradio-container code, | |
| .gradio-container .prose code { | |
| background: #171717 !important; | |
| color: #d4d4d5 !important; | |
| border: 1px solid #262626 !important; | |
| border-radius: 5px !important; | |
| padding: 2px 7px !important; | |
| font-family: 'JetBrains Mono', ui-monospace, monospace !important; | |
| font-size: 0.85em !important; | |
| } | |
| footer { display: none !important; } | |
| .gradio-container .prose a { color: """ + KREA_ACCENT + """; } | |
| """ | |
| with gr.Blocks(title="Krea 2") as demo: | |
| with gr.Column(elem_id="page"): | |
| gr.HTML( | |
| """ | |
| <header id="krea-header"> | |
| <div class="eyebrow">KREA · TEXT-TO-IMAGE</div> | |
| <h1>Krea 2</h1> | |
| <p class="subtitle">Generate images from natural language. Pick Raw for CFG-guided control or Turbo for fast, few-step results.</p> | |
| <div class="meta"> | |
| <div class="badges"> | |
| <span class="badge">Raw · CFG</span> | |
| <span class="badge">Turbo · few-step</span> | |
| </div> | |
| <div class="links"> | |
| <a href="https://www.krea.ai/blog/krea-2-technical-report" target="_blank" rel="noopener">Technical report ↗</a> | |
| <a href="https://github.com/krea-ai/krea-2" target="_blank" rel="noopener">GitHub ↗</a> | |
| </div> | |
| </div> | |
| </header> | |
| """ | |
| ) | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=5, elem_classes="panel"): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| lines=4, | |
| placeholder=PLACEHOLDER, | |
| show_label=True, | |
| autofocus=True, | |
| ) | |
| model = gr.Radio(["Turbo", "Raw"], value="Turbo", label="Model") | |
| run = gr.Button("Generate", variant="primary", elem_id="generate-btn") | |
| with gr.Accordion("Prompting tips", open=False): | |
| gr.Markdown(PROMPT_TIPS) | |
| resolution = gr.Radio( | |
| list(RESOLUTIONS.keys()), | |
| value="Square · 1024", | |
| label="Resolution", | |
| ) | |
| prompt_strength = gr.Slider( | |
| PROMPT_STRENGTH_MIN, | |
| PROMPT_STRENGTH_MAX, | |
| value=PROMPT_STRENGTH_DEFAULT, | |
| step=0.05, | |
| label="Prompt strength", | |
| info="How closely the image follows your prompt. 1.0 is the default; higher values amplify and rebalance conditioning to counteract the safety filter (recommended ~2-4).", | |
| ) | |
| with gr.Accordion("Advanced", open=False): | |
| negative_prompt = gr.Textbox( | |
| label="Negative prompt", | |
| lines=1, | |
| interactive=False, | |
| info="Available with Raw, where guidance is above 0.", | |
| ) | |
| steps = gr.Slider(1, 50, value=8, step=1, label="Steps") | |
| guidance = gr.Slider(0.0, 10.0, value=0.0, step=0.1, label="Guidance scale") | |
| 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") | |
| with gr.Row(): | |
| seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed") | |
| randomize = gr.Checkbox(value=True, label="Randomize seed") | |
| with gr.Column(scale=6, elem_classes="panel"): | |
| output = gr.Image(label="Result", format="png", elem_id="result-image") | |
| if GALLERY_ITEMS: | |
| gallery = gr.Gallery( | |
| value=GALLERY_ITEMS, | |
| label="Example prompts", | |
| columns=4, | |
| height="auto", | |
| object_fit="cover", | |
| allow_preview=False, | |
| elem_id="examples-gallery", | |
| ) | |
| def use_example(evt: gr.SelectData): | |
| # Clicking a sample loads its prompt into the box, ready to run. | |
| return GALLERY_PROMPTS[evt.index] | |
| gallery.select(use_example, None, prompt) | |
| else: | |
| # No bundled sample images present; show the prompts as text. | |
| gr.Examples( | |
| fn=generate, | |
| examples=EXAMPLE_PROMPTS, | |
| inputs=[prompt], | |
| outputs=[output, seed], | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Example prompts", | |
| examples_per_page=4, | |
| ) | |
| model.change(on_model_change, model, [steps, guidance, negative_prompt]) | |
| resolution.change(on_resolution_change, resolution, [width, height]) | |
| inputs = [prompt, negative_prompt, model, steps, guidance, width, height, seed, randomize, prompt_strength] | |
| run.click(generate, inputs, [output, seed]) | |
| prompt.submit(generate, inputs, [output, seed]) | |
| demo.launch(theme=theme, css=CSS) |