"""
Krea 2 Image Reference — training-free style transfer via Untwisting RoPE.

Renders a text prompt's *content* in a reference image's *style*, with no
training and no LoRA. Port of https://github.com/BigStationW/ComfyUi-Untwisting-RoPE
("Untwisting RoPE: Frequency Control for Shared Attention in DiTs",
https://arxiv.org/abs/2602.05013) to the 🧨 diffusers `Krea2Pipeline`.

Live demo: https://huggingface.co/spaces/multimodalart/Krea-2-Image-Reference

Setup
-----
    pip install "transformers>=4.57.0" accelerate sentencepiece \
        git+https://github.com/huggingface/diffusers.git
    # download the module that lives next to this file (from the gist or the Space):
    #   https://huggingface.co/spaces/multimodalart/Krea-2-Image-Reference/raw/main/krea2_untwist.py

Needs a CUDA GPU with enough memory for Krea 2 Turbo in bf16 (~26 GB) — the
style transfer runs a [target, reference] cross-batch, so budget accordingly.
"""

import torch
from diffusers import Krea2Pipeline
from diffusers.utils import load_image

# `krea2_untwist.py` must be importable (same folder as this script).
from krea2_untwist import style_transfer

pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", torch_dtype=torch.bfloat16)
pipe.to("cuda")

# The reference supplies the STYLE (palette / texture / rendering); the prompt
# supplies the CONTENT (composition / subject).
reference = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
)

image = style_transfer(
    pipe,
    prompt="a red fox sitting in a snowy forest, soft winter light",
    reference_image=reference,
    height=1024,
    width=1024,
    num_inference_steps=8,          # Krea 2 Turbo is few-step, guidance-free
    # --- style controls (these are the tuned defaults from the demo) ---
    beta=2.25,                      # sharpness of the frequency curve
    low_scale_start=1.0,            # low-freq (style) scale eases in ...
    low_scale_end=2.75,             # ... to `style_strength` at the last steps
    high_scale_start=1.0,           # high-freq (structure) scale decays ...
    high_scale_end=0.0,             # ... to zero, so composition follows the prompt
    adain_strength=0.75,            # match color/contrast statistics to the reference
    blocks=(7, 27),                 # skip early blocks so the prompt keeps its layout
    generator=torch.Generator("cuda").manual_seed(0),
)
image.save("krea2_image_reference.png")
print("saved krea2_image_reference.png")

# Tuning tips
# -----------
# * Weak effect?  Raise low_scale_end (style strength) toward 3.0, raise
#   adain_strength, or lower the first block toward 0.
# * Reference bleeding into composition?  Lower low_scale_end, keep
#   high_scale_end at 0.0, and keep the first block >= 7.