Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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( | |
| """ | |
| <header id="hdr"> | |
| <div class="eyebrow">KREA 2 · IMAGE REFERENCE</div> | |
| <h1>Krea 2 Image Reference</h1> | |
| <p class="subtitle">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.</p> | |
| </header> | |
| """ | |
| ) | |
| 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() | |