import os import random import gradio as gr import spaces import torch from huggingface_hub import login if os.environ.get("HF_TOKEN"): login(token=os.environ["HF_TOKEN"]) from diffusers import Krea2Pipeline from krea2_untwist import style_transfer DTYPE = torch.bfloat16 TURBO_REPO = "krea/Krea-2-Turbo" MAX_SEED = 2**31 - 1 # Krea 2 Turbo (few-step, guidance-free) is the fast path — ideal for the # interactive image-reference loop. Untwisting-RoPE swaps the attention # processors at run time, which is incompatible with AOTI-compiled blocks, so # this Space runs eager on purpose. pipe = Krea2Pipeline.from_pretrained(TURBO_REPO, torch_dtype=DTYPE) pipe.to("cuda") RESOLUTIONS = { "Square · 1024": (1024, 1024), "Portrait · 1024": (832, 1216), "Landscape · 1024": (1216, 832), } def _duration(*args, **kwargs): steps = args[2] if len(args) > 2 else 8 # matches run()'s positional order return int(int(steps) * 6 + 40) # cross-batch (target+reference) ~2x per step @spaces.GPU(duration=_duration, size="xlarge") def run( reference_image, prompt, steps=8, style_strength=2.75, structure=1.0, beta=2.25, adain_strength=0.75, block_start=7, block_end=27, resolution="Square · 1024", seed=0, randomize=True, progress=gr.Progress(track_tqdm=True), ): if reference_image is None: raise gr.Error("Upload a reference image to guide the style.") if not prompt or not prompt.strip(): raise gr.Error("Enter a prompt describing what to generate.") if randomize: seed = random.randint(0, MAX_SEED) seed = int(seed) width, height = RESOLUTIONS[resolution] generator = torch.Generator("cuda").manual_seed(seed) try: image = style_transfer( pipe, prompt=prompt, reference_image=reference_image, height=height, width=width, num_inference_steps=int(steps), guidance_scale=0.0, beta=float(beta), # High-frequency (structure) contribution decays to zero over # denoising so composition follows the prompt, not the reference. high_scale_start=float(structure), high_scale_end=0.0, # Low-frequency (style) contribution eases in to `style_strength`. low_scale_start=1.0, low_scale_end=float(style_strength), adain_strength=float(adain_strength), blocks=(int(block_start), int(block_end)), generator=generator, ) except torch.cuda.OutOfMemoryError as exc: torch.cuda.empty_cache() raise gr.Error("Ran out of GPU memory. Try a 1024 resolution.") from exc return image, seed ABOUT = """\ **Krea 2 Image Reference** does training-free **style transfer**: give a reference image and a prompt, and Krea 2 renders your prompt's *content* in the reference's *style*. It ports [Untwisting RoPE: Frequency Control for Shared Attention in DiTs](https://arxiv.org/abs/2602.05013) ([ComfyUI original](https://github.com/BigStationW/ComfyUi-Untwisting-RoPE)) to 🧨 diffusers. At every denoising step the model runs on a `[target, reference]` pair. Inside each attention block, after rotary position embedding, the reference image's **keys** are rescaled per frequency band and the target attends to those reference keys/values. The high-frequency (fine-structure) contribution decays to zero across the steps, so the reference guides *style* — palette, texture, rendering — while composition follows your prompt. **Controls** - **Style strength** — how strongly the reference style comes through. - **Structure** — how much of the reference's layout leaks in (keep low for style-only; raise to echo the reference composition). - **beta** — sharpness of the frequency curve. - **AdaIN** — matches color/contrast statistics to the reference. - **Blocks** — which transformer blocks are affected (skipping early blocks keeps the prompt's composition intact). No training, no LoRA. Runs eager (the attention-processor swap is incompatible with AOTI-compiled blocks). **Use it in code** — `krea2_untwist.py` is a standalone module that works with plain diffusers. Full example + module in this [gist](https://gist.github.com/apolinario/9ecc9e0efffbbf133fe997b7181b6cfa); then `from krea2_untwist import style_transfer`. """ 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_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, slider_color=KREA_ACCENT, ) CSS = """ .gradio-container { background: #000 !important; } #page { max-width: 1120px; margin: 0 auto; padding: 4px 8px 32px; } #hdr { padding: 30px 6px 20px; border-bottom: 1px solid #1a1a1a; margin-bottom: 20px; } #hdr .eyebrow { font-family: 'JetBrains Mono', monospace; font-size: 11px; letter-spacing: 0.24em; text-transform: uppercase; color: #737373; } #hdr h1 { font-size: 38px; font-weight: 600; letter-spacing: -0.025em; margin: 10px 0 6px; color: #fff; } #hdr .subtitle { font-size: 15px; line-height: 1.5; color: #a3a3a3; margin: 0; max-width: 62ch; } #go { font-weight: 600; } #result { min-height: 420px; border-radius: 10px; overflow: hidden; } footer { display: none !important; } .gradio-container .prose a { color: """ + KREA_ACCENT + """; } """ with gr.Blocks(title="Krea 2 Image Reference", theme=theme, css=CSS) as demo: with gr.Column(elem_id="page"): gr.HTML( """
KREA 2 · IMAGE REFERENCE

Krea 2 Image Reference

Give a reference image and a prompt. Krea 2 renders your prompt in the reference's style — training-free, via RoPE frequency control over shared attention.

""" ) with gr.Row(equal_height=False): with gr.Column(scale=5): reference_image = gr.Image(label="Reference image (style)", type="pil", height=280) prompt = gr.Textbox( label="Prompt (content)", lines=3, placeholder="a lighthouse on a cliff at sunset, crashing waves", autofocus=True, ) go = gr.Button("Generate", variant="primary", elem_id="go") style_strength = gr.Slider(1.0, 3.0, value=2.75, step=0.05, label="Style strength") resolution = gr.Radio(list(RESOLUTIONS.keys()), value="Square · 1024", label="Resolution") with gr.Accordion("Style controls", open=False): structure = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Structure (reference layout carry-over)") beta = gr.Slider(1.0, 8.0, value=2.25, step=0.1, label="beta · curve sharpness") adain_strength = gr.Slider(0.0, 1.0, value=0.75, step=0.05, label="AdaIN · color/contrast match") with gr.Accordion("Advanced", open=False): steps = gr.Slider(4, 16, value=8, step=1, label="Steps") with gr.Row(): block_start = gr.Slider(0, 27, value=7, step=1, label="First affected block") block_end = gr.Slider(0, 27, value=27, step=1, label="Last affected block") 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): output = gr.Image(label="Result", format="png", elem_id="result") out_seed = gr.Number(label="Seed used", interactive=False) with gr.Accordion("How it works", open=False): gr.Markdown(ABOUT) inputs = [reference_image, prompt, steps, style_strength, structure, beta, adain_strength, block_start, block_end, resolution, seed, randomize] go.click(run, inputs, [output, out_seed]) prompt.submit(run, inputs, [output, out_seed]) demo.launch()