app.py

"""Krea 2 Turbo – OpenPose ControlNet LoRA demo.

Two ways to drive the pose:

* **Pose Image** tab — upload an already-rendered OpenPose/DWPose skeleton
  image; it is used as the control input as-is.
* **Regular Image** tab — upload a normal photo; a DWPose skeleton is
  auto-extracted and shown in an interactive editor where individual joints
  can be dragged. The (possibly edited) skeleton is what actually conditions
  generation.

The pose map is fed to the Krea 2 model via the Ostris Edit reference-image
conditioning path (Qwen3-VL vision tokens + clean VAE reference latents at
t=0). To make the model actually follow the pose, the generation resolution is
matched to the pose image's aspect ratio/size (snapped to a multiple of 16) —
the pose latents sit on their own rotary index grid starting at (0,0), so they
only line up with the output when both grids share the same height/width.

Base model: krea/Krea-2-Turbo (gated, auto-approve)
LoRA:       thedeoxen/Krea-2-pose-controlnet
Interactive skeleton editor inspired by
https://huggingface.co/spaces/linoyts/Flux-2-control-pose
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / diffusers / transformers
import torch
import gradio as gr
import numpy as np
import cv2
from PIL import Image
from diffusers import DiffusionPipeline

# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda") per ZeroGPU rules)
# ---------------------------------------------------------------------------

BASE_MODEL = "krea/Krea-2-Turbo"
LORA_REPO = "thedeoxen/Krea-2-pose-controlnet"
LORA_WEIGHT = "krea2_turbo_openpose_controlnet.safetensors"

_token = os.environ.get("HF_TOKEN")

pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL,
    custom_pipeline="ostris/Krea2OstrisEdit",
    torch_dtype=torch.bfloat16,
    token=_token,
    trust_remote_code=True,
)
pipe.to("cuda")
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHT, token=_token)
pipe.transformer.set_adapters("default", weights=1.0)

# ---------------------------------------------------------------------------
# Pose extraction (DWPose via controlnet_aux) — runs on CPU, module scope
# ---------------------------------------------------------------------------

from controlnet_aux import OpenposeDetector

pose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators")

# ---------------------------------------------------------------------------
# Interactive pose-skeleton editor component + keypoint <-> image helpers.
# Keypoint layout (adapted from linoyts/Flux-2-control-pose):
#   18 body + 70 face + 21 left hand + 21 right hand = 130 total.
# Each keypoint is {"x", "y", "vis"} with x/y normalized to [0, 1].
# ---------------------------------------------------------------------------

TOTAL_KEYPOINTS = 130
BODY_COUNT = 18
FACE_COUNT = 70
HAND_COUNT = 21
FACE_START = 18
LEFT_HAND_START = 88
RIGHT_HAND_START = 109


def _snap16(v):
    """Round a dimension up to the nearest multiple of 16, clamped to [512, 1024]."""
    v = int(round(v / 16.0) * 16)
    return max(512, min(1024, v))


# Exact replica of the ComfyUI reference workflow's sizing node
# (FluxKontextImageScale, node 28). In the workflow the DWPose skeleton is run
# through FluxKontextImageScale, and that single scaled image drives BOTH:
#   * node 35 image1 -> the VAE reference latents, and
#   * node 6 EmptyLatentImage (via GetImageSize) -> the output canvas.
# So the reference-latent grid and the output grid are guaranteed identical.
#
# The previous "fit under 1 MP, snap /16, never upscale" heuristic diverged:
# small pose images (e.g. the 512x768 editor canvas or typical skeleton PNGs)
# were left at ~0.4 MP, whereas FluxKontextImageScale always normalizes to a
# ~1 MP preferred bucket. Krea-2-Turbo is trained at ~1 MP, so running the
# pose reference + generation at a third of that resolution destroyed both
# generation quality and pose fidelity (the model saw a coarse skeleton on a
# small grid). This replicates FluxKontextImageScale precisely.
PREFERRED_KONTEXT_RESOLUTIONS = [
    (672, 1568), (688, 1504), (720, 1456), (752, 1392), (800, 1328),
    (832, 1248), (880, 1184), (944, 1104), (1024, 1024), (1104, 944),
    (1184, 880), (1248, 832), (1328, 800), (1392, 752), (1456, 720),
    (1504, 688), (1568, 672),
]


def _flux_kontext_scale(pose_image):
    """Scale a pose image exactly like ComfyUI's FluxKontextImageScale node.

    Picks the preferred (~1 MP) resolution bucket with the nearest aspect
    ratio, center-crops to that aspect, then Lanczos-resizes to the exact
    bucket size. Returns (scaled_rgb_image, width, height). Both bucket dims
    are multiples of 16, so the Krea-2 latent grid patchifies cleanly and the
    output grid lines up with the pose reference-latent grid.
    """
    pose_image = pose_image.convert("RGB")
    w, h = pose_image.size
    aspect = w / h
    _, bw, bh = min(
        (abs(aspect - tw / th), tw, th) for tw, th in PREFERRED_KONTEXT_RESOLUTIONS
    )

    # Center-crop the source to the target aspect (crop only, never pad),
    # matching comfy.utils.common_upscale(crop="center").
    old_aspect = w / h
    new_aspect = bw / bh
    x = y = 0
    if old_aspect > new_aspect:
        x = round((w - w * (new_aspect / old_aspect)) / 2)
    elif old_aspect < new_aspect:
        y = round((h - h * (old_aspect / new_aspect)) / 2)
    if x > 0 or y > 0:
        pose_image = pose_image.crop((x, y, w - x, h - y))

    # Lanczos resize to the exact bucket size (PIL LANCZOS, as ComfyUI does).
    pose_image = pose_image.resize((bw, bh), resample=Image.LANCZOS)
    return pose_image, bw, bh


class InteractivePoseSkeleton(gr.HTML):
    """Interactive pose skeleton editor with draggable keypoints."""

    def __init__(self, value=None, width=512, height=768, **kwargs):
        # Body skeleton connections (indices 0-17)
        connections = [
            [0, 1], [1, 2], [2, 3], [3, 4],
            [1, 5], [5, 6], [6, 7],
            [1, 8], [8, 9], [9, 10],
            [1, 11], [11, 12], [12, 13],
            [0, 14], [14, 16], [0, 15], [15, 17],
        ]

        face_start = 18
        for i in range(16):
            connections.append([face_start + i, face_start + i + 1])
        for i in range(17, 21):
            connections.append([face_start + i, face_start + i + 1])
        for i in range(22, 26):
            connections.append([face_start + i, face_start + i + 1])
        for i in range(27, 30):
            connections.append([face_start + i, face_start + i + 1])
        for i in range(31, 35):
            connections.append([face_start + i, face_start + i + 1])
        for i in range(36, 41):
            connections.append([face_start + i, face_start + i + 1])
        connections.append([face_start + 41, face_start + 36])
        for i in range(42, 47):
            connections.append([face_start + i, face_start + i + 1])
        connections.append([face_start + 47, face_start + 42])
        for i in range(48, 59):
            connections.append([face_start + i, face_start + i + 1])
        connections.append([face_start + 59, face_start + 48])
        for i in range(60, 67):
            connections.append([face_start + i, face_start + i + 1])
        connections.append([face_start + 67, face_start + 60])

        for hand_start in [88, 109]:
            for i in range(4):
                connections.append([hand_start + i, hand_start + i + 1])
            for i in range(5, 8):
                connections.append([hand_start + i, hand_start + i + 1])
            for i in range(9, 12):
                connections.append([hand_start + i, hand_start + i + 1])
            for i in range(13, 16):
                connections.append([hand_start + i, hand_start + i + 1])
            for i in range(17, 20):
                connections.append([hand_start + i, hand_start + i + 1])
            connections.append([hand_start + 0, hand_start + 5])
            connections.append([hand_start + 0, hand_start + 9])
            connections.append([hand_start + 0, hand_start + 13])
            connections.append([hand_start + 0, hand_start + 17])

        connections.append([7, 88])
        connections.append([4, 109])

        html_template = """
        <div class="skeleton-editor">
            <canvas id="skeleton-canvas"></canvas>
            <button id="reset-btn" class="btn-reset">Reset</button>
        </div>
        """

        css_template = """
        .skeleton-editor {
            position: relative;
            display: flex;
            flex-direction: column;
            align-items: center;
            padding: 20px;
            background: linear-gradient(135deg, rgba(0,0,0,0.95) 0%, rgba(20,20,30,0.95) 100%);
            border-radius: 12px;
            box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
        }
        #skeleton-canvas {
            border: 3px solid transparent;
            border-radius: 10px;
            cursor: crosshair;
            background: #000000;
            width: auto;
            height: ${height}px;
            max-width: 100%;
            aspect-ratio: ${width} / ${height};
            display: block;
            box-shadow:
                0 0 20px rgba(0, 255, 255, 0.3),
                0 0 40px rgba(0, 255, 255, 0.2),
                inset 0 0 20px rgba(0, 255, 255, 0.1);
            transition: box-shadow 0.3s ease;
        }
        #skeleton-canvas:hover {
            box-shadow:
                0 0 30px rgba(0, 255, 255, 0.5),
                0 0 60px rgba(0, 255, 255, 0.3),
                inset 0 0 30px rgba(0, 255, 255, 0.2);
        }
        .btn-reset {
            position: absolute;
            top: 28px;
            right: 28px;
            padding: 8px 16px;
            border: 1px solid rgba(255, 255, 255, 0.2);
            border-radius: 6px;
            cursor: pointer;
            font-size: 12px;
            font-weight: 600;
            background: rgba(0, 0, 0, 0.6);
            color: #ffffff;
            backdrop-filter: blur(10px);
            transition: all 0.2s ease;
            z-index: 10;
            text-transform: uppercase;
            letter-spacing: 0.5px;
        }
        .btn-reset:hover {
            background: rgba(255, 255, 255, 0.1);
            border-color: rgba(0, 255, 255, 0.5);
            box-shadow: 0 0 10px rgba(0, 255, 255, 0.3);
        }
        """

        js_on_load = f"""
        const canvas = element.querySelector('#skeleton-canvas');
        const ctx = canvas.getContext('2d');
        const resetBtn = element.querySelector('#reset-btn');

        canvas.width = {width};
        canvas.height = {height};

        const connections = {connections};

        const bodyLimbColors = [
            '#FF0000', '#FF5500', '#FFAA00', '#FFFF00',
            '#AAFF00', '#55FF00', '#00FF00', '#00FF55',
            '#00FFAA', '#00FFFF', '#00AAFF', '#0055FF',
            '#0000FF', '#5500FF', '#AA00FF', '#FF00FF',
            '#FF00AA', '#FF0055'
        ];
        const bodyKeypointColors = bodyLimbColors;
        const handFingerColors = ['#FF0000', '#FF8800', '#FFFF00', '#00FF00', '#0000FF'];

        const connectionColors = {{
            '1-8': 9, '1-2': 8, '1-5': 7,
            '2-3': 10, '3-4': 11, '5-6': 6, '6-7': 5,
            '8-9': 12, '9-10': 13, '1-11': 4,
            '11-12': 3, '12-13': 2, '1-0': 1,
            '0-14': 16, '14-16': 18, '0-15': 15, '15-17': 17
        }};

        let keypoints = props.value ? JSON.parse(JSON.stringify(props.value)) : [];
        let originalKeypoints = props.value ? JSON.parse(JSON.stringify(props.value)) : [];
        let draggingIndex = -1;
        let hoveredIndex = -1;
        let prevDragX = -1;
        let prevDragY = -1;

        const childrenMap = {{
            1:  [0, 2, 5, 8, 11],
            0:  [14, 15],
            14: [16],
            15: [17],
            2:  [3],
            3:  [4],
            5:  [6],
            6:  [7],
            8:  [9],
            9:  [10],
            11: [12],
            12: [13],
        }};

        const attachedGroups = {{
            0:  Array.from({{length: 70}}, (_, i) => 18 + i),
            7:  Array.from({{length: 21}}, (_, i) => 88 + i),
            4:  Array.from({{length: 21}}, (_, i) => 109 + i),
        }};

        function getDescendants(idx) {{
            const result = [];
            const stack = [idx];
            const visited = new Set();
            visited.add(idx);
            while (stack.length > 0) {{
                const current = stack.pop();
                const bodyKids = childrenMap[current] || [];
                for (const kid of bodyKids) {{
                    if (!visited.has(kid)) {{
                        visited.add(kid);
                        result.push(kid);
                        stack.push(kid);
                    }}
                }}
                const attached = attachedGroups[current] || [];
                for (const ai of attached) {{
                    if (!visited.has(ai)) {{
                        visited.add(ai);
                        result.push(ai);
                    }}
                }}
            }}
            return result;
        }}

        function drawSkeleton() {{
            ctx.clearRect(0, 0, canvas.width, canvas.height);
            if (!keypoints || keypoints.length === 0) {{ return; }}

            ctx.lineCap = 'round';
            const bodyConnections = connections.filter(([i, j]) => i < 18 && j < 18);
            bodyConnections.forEach(([i, j], connIdx) => {{
                if (i < keypoints.length && j < keypoints.length) {{
                    const kp_i = keypoints[i];
                    const kp_j = keypoints[j];
                    if (kp_i && kp_j && kp_i.vis > 0.1 && kp_j.vis > 0.1) {{
                        const x1 = kp_i.x * canvas.width;
                        const y1 = kp_i.y * canvas.height;
                        const x2 = kp_j.x * canvas.width;
                        const y2 = kp_j.y * canvas.height;
                        const key = `${{i}}-${{j}}`;
                        const reverseKey = `${{j}}-${{i}}`;
                        const colorIdx = connectionColors[key] || connectionColors[reverseKey] || connIdx;
                        const color = bodyLimbColors[colorIdx % bodyLimbColors.length];
                        ctx.lineWidth = 6;
                        ctx.shadowBlur = 15;
                        ctx.shadowColor = color;
                        ctx.strokeStyle = color;
                        ctx.beginPath();
                        ctx.moveTo(x1, y1);
                        ctx.lineTo(x2, y2);
                        ctx.stroke();
                    }}
                }}
            }});
            ctx.shadowBlur = 0;

            const bridgeConns = [[7, 88], [4, 109]];
            bridgeConns.forEach(([i, j]) => {{
                if (i < keypoints.length && j < keypoints.length) {{
                    const kp_i = keypoints[i];
                    const kp_j = keypoints[j];
                    if (kp_i && kp_j && kp_i.vis > 0.1 && kp_j.vis > 0.1) {{
                        const x1 = kp_i.x * canvas.width;
                        const y1 = kp_i.y * canvas.height;
                        const x2 = kp_j.x * canvas.width;
                        const y2 = kp_j.y * canvas.height;
                        const color = bodyKeypointColors[i % bodyKeypointColors.length];
                        ctx.lineWidth = 3;
                        ctx.setLineDash([4, 4]);
                        ctx.strokeStyle = color;
                        ctx.beginPath();
                        ctx.moveTo(x1, y1);
                        ctx.lineTo(x2, y2);
                        ctx.stroke();
                        ctx.setLineDash([]);
                    }}
                }}
            }});

            const faceConns = connections.filter(([i, j]) => i >= 18 && i < 88 && j >= 18 && j < 88);
            faceConns.forEach(([i, j]) => {{
                if (i < keypoints.length && j < keypoints.length) {{
                    const kp_i = keypoints[i];
                    const kp_j = keypoints[j];
                    if (kp_i && kp_j && kp_i.vis > 0.1 && kp_j.vis > 0.1) {{
                        ctx.lineWidth = 1;
                        ctx.strokeStyle = 'rgba(255, 255, 255, 0.4)';
                        ctx.beginPath();
                        ctx.moveTo(kp_i.x * canvas.width, kp_i.y * canvas.height);
                        ctx.lineTo(kp_j.x * canvas.width, kp_j.y * canvas.height);
                        ctx.stroke();
                    }}
                }}
            }});

            [88, 109].forEach(handStart => {{
                const fingerConns = [];
                for (let fi = 0; fi < 4; fi++) fingerConns.push([handStart + fi, handStart + fi + 1, 0]);
                for (let fi = 5; fi < 8; fi++) fingerConns.push([handStart + fi, handStart + fi + 1, 1]);
                for (let fi = 9; fi < 12; fi++) fingerConns.push([handStart + fi, handStart + fi + 1, 2]);
                for (let fi = 13; fi < 16; fi++) fingerConns.push([handStart + fi, handStart + fi + 1, 3]);
                for (let fi = 17; fi < 20; fi++) fingerConns.push([handStart + fi, handStart + fi + 1, 4]);
                fingerConns.push([handStart + 0, handStart + 5, 1]);
                fingerConns.push([handStart + 0, handStart + 9, 2]);
                fingerConns.push([handStart + 0, handStart + 13, 3]);
                fingerConns.push([handStart + 0, handStart + 17, 4]);
                fingerConns.forEach(([i, j, fingerIdx]) => {{
                    if (i < keypoints.length && j < keypoints.length) {{
                        const kp_i = keypoints[i];
                        const kp_j = keypoints[j];
                        if (kp_i && kp_j && kp_i.vis > 0.1 && kp_j.vis > 0.1) {{
                            const color = handFingerColors[fingerIdx];
                            ctx.lineWidth = 2;
                            ctx.strokeStyle = color;
                            ctx.beginPath();
                            ctx.moveTo(kp_i.x * canvas.width, kp_i.y * canvas.height);
                            ctx.lineTo(kp_j.x * canvas.width, kp_j.y * canvas.height);
                            ctx.stroke();
                        }}
                    }}
                }});
            }});
            ctx.shadowBlur = 0;

            for (let index = 0; index < Math.min(18, keypoints.length); index++) {{
                const kp = keypoints[index];
                if (kp && kp.vis > 0.1) {{
                    const x = kp.x * canvas.width;
                    const y = kp.y * canvas.height;
                    const isActive = index === draggingIndex || index === hoveredIndex;
                    const color = bodyKeypointColors[index % bodyKeypointColors.length];
                    const baseRadius = isActive ? 14 : 10;
                    ctx.shadowBlur = isActive ? 25 : 12;
                    ctx.shadowColor = color;
                    ctx.fillStyle = '#ffffff';
                    ctx.beginPath();
                    ctx.arc(x, y, baseRadius + 2, 0, 2 * Math.PI);
                    ctx.fill();
                    ctx.fillStyle = color;
                    ctx.beginPath();
                    ctx.arc(x, y, baseRadius, 0, 2 * Math.PI);
                    ctx.fill();
                    ctx.shadowBlur = 0;
                }}
            }}

            for (let index = 18; index < Math.min(88, keypoints.length); index++) {{
                const kp = keypoints[index];
                if (kp && kp.vis > 0.1) {{
                    const isActive = index === draggingIndex || index === hoveredIndex;
                    ctx.fillStyle = isActive ? '#00ffff' : '#ffffff';
                    ctx.beginPath();
                    ctx.arc(kp.x * canvas.width, kp.y * canvas.height, isActive ? 5 : 3, 0, 2 * Math.PI);
                    ctx.fill();
                }}
            }}

            [88, 109].forEach(handStart => {{
                for (let hi = 0; hi < 21; hi++) {{
                    const index = handStart + hi;
                    if (index < keypoints.length) {{
                        const kp = keypoints[index];
                        if (kp && kp.vis > 0.1) {{
                            let fingerIdx;
                            if (hi <= 4) fingerIdx = 0;
                            else if (hi <= 8) fingerIdx = 1;
                            else if (hi <= 12) fingerIdx = 2;
                            else if (hi <= 16) fingerIdx = 3;
                            else fingerIdx = 4;
                            const color = handFingerColors[fingerIdx];
                            const isActive = index === draggingIndex || index === hoveredIndex;
                            ctx.fillStyle = color;
                            ctx.beginPath();
                            ctx.arc(kp.x * canvas.width, kp.y * canvas.height, isActive ? 6 : 4, 0, 2 * Math.PI);
                            ctx.fill();
                        }}
                    }}
                }}
            }});
        }}

        function getMousePos(e) {{
            const rect = canvas.getBoundingClientRect();
            const scaleX = canvas.width / rect.width;
            const scaleY = canvas.height / rect.height;
            return {{
                x: (e.clientX - rect.left) * scaleX / canvas.width,
                y: (e.clientY - rect.top) * scaleY / canvas.height
            }};
        }}

        function findNearestKeypoint(pos) {{
            let nearest = -1;
            let minDist = Infinity;
            for (let index = 0; index < Math.min(18, keypoints.length); index++) {{
                const kp = keypoints[index];
                if (kp && kp.vis > 0.1) {{
                    const dx = kp.x - pos.x;
                    const dy = kp.y - pos.y;
                    const dist = Math.sqrt(dx * dx + dy * dy);
                    if (dist < 0.04 && dist < minDist) {{ minDist = dist; nearest = index; }}
                }}
            }}
            if (nearest === -1) {{
                for (let index = 18; index < keypoints.length; index++) {{
                    const kp = keypoints[index];
                    if (kp && kp.vis > 0.1) {{
                        const dx = kp.x - pos.x;
                        const dy = kp.y - pos.y;
                        const dist = Math.sqrt(dx * dx + dy * dy);
                        if (dist < 0.025 && dist < minDist) {{ minDist = dist; nearest = index; }}
                    }}
                }}
            }}
            return nearest;
        }}

        function commit() {{
            props.value = keypoints.map(kp => ({{
                x: parseFloat(kp.x) || 0,
                y: parseFloat(kp.y) || 0,
                vis: parseFloat(kp.vis) || 0
            }}));
            trigger('change');
        }}

        canvas.addEventListener('mousedown', (e) => {{
            const pos = getMousePos(e);
            draggingIndex = findNearestKeypoint(pos);
            if (draggingIndex !== -1) {{
                prevDragX = pos.x;
                prevDragY = pos.y;
                canvas.style.cursor = 'grabbing';
                drawSkeleton();
            }}
        }});

        canvas.addEventListener('mousemove', (e) => {{
            const pos = getMousePos(e);
            if (draggingIndex !== -1) {{
                const dx = pos.x - prevDragX;
                const dy = pos.y - prevDragY;
                prevDragX = pos.x;
                prevDragY = pos.y;
                keypoints[draggingIndex].x = Math.max(0, Math.min(1, keypoints[draggingIndex].x + dx));
                keypoints[draggingIndex].y = Math.max(0, Math.min(1, keypoints[draggingIndex].y + dy));
                if (draggingIndex < 18) {{
                    const descendants = getDescendants(draggingIndex);
                    for (const di of descendants) {{
                        if (di < keypoints.length && keypoints[di] && keypoints[di].vis > 0.01) {{
                            keypoints[di].x = Math.max(0, Math.min(1, keypoints[di].x + dx));
                            keypoints[di].y = Math.max(0, Math.min(1, keypoints[di].y + dy));
                        }}
                    }}
                }} else if (draggingIndex === 88 || draggingIndex === 109) {{
                    const handStart = draggingIndex;
                    for (let hi = 1; hi < 21; hi++) {{
                        const gi = handStart + hi;
                        if (gi < keypoints.length && keypoints[gi] && keypoints[gi].vis > 0.01) {{
                            keypoints[gi].x = Math.max(0, Math.min(1, keypoints[gi].x + dx));
                            keypoints[gi].y = Math.max(0, Math.min(1, keypoints[gi].y + dy));
                        }}
                    }}
                }}
                drawSkeleton();
            }} else {{
                const newHovered = findNearestKeypoint(pos);
                if (newHovered !== hoveredIndex) {{
                    hoveredIndex = newHovered;
                    canvas.style.cursor = hoveredIndex !== -1 ? 'grab' : 'crosshair';
                    drawSkeleton();
                }}
            }}
        }});

        canvas.addEventListener('mouseup', () => {{
            if (draggingIndex !== -1) {{
                draggingIndex = -1;
                prevDragX = -1;
                prevDragY = -1;
                canvas.style.cursor = hoveredIndex !== -1 ? 'grab' : 'crosshair';
                commit();
            }}
        }});

        canvas.addEventListener('mouseleave', () => {{
            if (draggingIndex !== -1) {{ commit(); }}
            draggingIndex = -1;
            prevDragX = -1;
            prevDragY = -1;
            hoveredIndex = -1;
            canvas.style.cursor = 'crosshair';
            drawSkeleton();
        }});

        resetBtn.addEventListener('click', () => {{
            keypoints = originalKeypoints.map(kp => ({{...kp}}));
            commit();
            drawSkeleton();
        }});

        let lastValue = null;
        const checkUpdates = () => {{
            const currentValue = JSON.stringify(props.value);
            if (currentValue !== lastValue && props.value && Array.isArray(props.value)) {{
                lastValue = currentValue;
                keypoints = props.value.map(kp => ({{
                    x: parseFloat(kp.x) || 0,
                    y: parseFloat(kp.y) || 0,
                    vis: parseFloat(kp.vis) || 0
                }}));
                originalKeypoints = keypoints.map(kp => ({{...kp}}));
                drawSkeleton();
            }}
            requestAnimationFrame(checkUpdates);
        }};
        checkUpdates();
        drawSkeleton();
        """

        super().__init__(
            value=value,
            width=width,
            height=height,
            html_template=html_template,
            css_template=css_template,
            js_on_load=js_on_load,
            **kwargs,
        )

    def api_info(self):
        return {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "x": {"type": "number"},
                    "y": {"type": "number"},
                    "vis": {"type": "number"},
                },
            },
        }


# ---------------------------------------------------------------------------
# Keypoint extraction / rendering helpers
# ---------------------------------------------------------------------------

def _as_list(part):
    if part is None:
        return []
    if isinstance(part, list):
        return part
    if hasattr(part, "keypoints"):
        return part.keypoints or []
    return []


def _append_kp(out, kp):
    if kp is None or not hasattr(kp, "x") or not hasattr(kp, "y"):
        out.append({"x": 0.0, "y": 0.0, "vis": 0.0})
        return
    out.append({
        "x": float(kp.x),
        "y": float(kp.y),
        "vis": float(getattr(kp, "score", 1.0)),
    })


def extract_pose_keypoints(image):
    """Extract normalized OpenPose keypoints (body/face/hands) from a photo.

    Returns (keypoints_list, pose_preview_image) or (None, None) if no pose.
    """
    if image is None:
        return None, None

    pose_img = pose_detector(
        image,
        hand_and_face=True,
        include_body=True,
        include_hand=True,
        include_face=True,
        output_type="pil",
    )

    detected = pose_detector.detect_poses(
        np.array(image),
        include_hand=True,
        include_face=True,
    )
    if not detected or len(detected) == 0:
        gr.Warning("No pose detected in image. Please try another image.")
        return None, None

    pose_data = detected[0]
    keypoints = []

    body_list = _as_list(getattr(pose_data, "body", None))
    for i in range(BODY_COUNT):
        _append_kp(keypoints, body_list[i] if i < len(body_list) else None)

    face_list = _as_list(getattr(pose_data, "face", None))
    for i in range(FACE_COUNT):
        _append_kp(keypoints, face_list[i] if i < len(face_list) else None)

    lhand_list = _as_list(getattr(pose_data, "left_hand", None))
    for i in range(HAND_COUNT):
        _append_kp(keypoints, lhand_list[i] if i < len(lhand_list) else None)

    rhand_list = _as_list(getattr(pose_data, "right_hand", None))
    for i in range(HAND_COUNT):
        _append_kp(keypoints, rhand_list[i] if i < len(rhand_list) else None)

    if len(keypoints) < TOTAL_KEYPOINTS:
        keypoints.extend([{"x": 0.0, "y": 0.0, "vis": 0.0}] * (TOTAL_KEYPOINTS - len(keypoints)))
    else:
        keypoints = keypoints[:TOTAL_KEYPOINTS]

    # Sanitize non-body keypoints: OpenPose parks undetected joints at (0,0)
    # or image edges with spurious scores.
    for i in range(BODY_COUNT, len(keypoints)):
        kp = keypoints[i]
        if kp["vis"] > 0:
            x, y = kp["x"], kp["y"]
            if (x == 0.0 and y == 0.0) or x <= 0.001 or y <= 0.001 or x >= 0.999 or y >= 0.999:
                kp["vis"] = 0.0

    hand_ranges = [
        (LEFT_HAND_START, LEFT_HAND_START + HAND_COUNT, 7, 6),
        (RIGHT_HAND_START, RIGHT_HAND_START + HAND_COUNT, 4, 3),
    ]
    for hand_start, hand_end, wrist_idx, elbow_idx in hand_ranges:
        wrist = keypoints[wrist_idx]
        if wrist["vis"] < 0.1:
            for i in range(hand_start, min(hand_end, len(keypoints))):
                keypoints[i]["vis"] = 0.0
            continue
        elbow = keypoints[elbow_idx]
        if elbow["vis"] > 0.1:
            forearm = ((wrist["x"] - elbow["x"]) ** 2 + (wrist["y"] - elbow["y"]) ** 2) ** 0.5
            max_dist = max(forearm * 1.8, 0.08)
        else:
            max_dist = 0.15
        for i in range(hand_start, min(hand_end, len(keypoints))):
            kp = keypoints[i]
            if kp["vis"] > 0.1:
                dist = ((kp["x"] - wrist["x"]) ** 2 + (kp["y"] - wrist["y"]) ** 2) ** 0.5
                if dist > max_dist:
                    kp["vis"] = 0.0

    nose = keypoints[0]
    if nose["vis"] > 0.1:
        neck = keypoints[1]
        if neck["vis"] > 0.1:
            head = ((nose["x"] - neck["x"]) ** 2 + (nose["y"] - neck["y"]) ** 2) ** 0.5
            face_max = max(head * 2.5, 0.10)
        else:
            face_max = 0.15
        for i in range(FACE_START, min(FACE_START + FACE_COUNT, len(keypoints))):
            kp = keypoints[i]
            if kp["vis"] > 0.1:
                dist = ((kp["x"] - nose["x"]) ** 2 + (kp["y"] - nose["y"]) ** 2) ** 0.5
                if dist > face_max:
                    kp["vis"] = 0.0

    body_vis = len([kp for kp in keypoints[:18] if kp["vis"] > 0.1])
    if body_vis < 8:
        gr.Warning("Incomplete pose detected. Please try a clearer image.")
        return None, None

    return keypoints, pose_img


def keypoints_to_pose_image(keypoints, width, height):
    """Render normalized keypoints into an OpenPose-format skeleton on black."""
    canvas = np.zeros((height, width, 3), dtype=np.uint8)
    if not keypoints or len(keypoints) == 0:
        return Image.fromarray(canvas)

    pose_colors = [
        [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0],
        [170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85],
        [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
        [0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255],
        [255, 0, 170], [255, 0, 85], [255, 0, 0],
    ]
    hand_finger_colors = [
        [255, 0, 0], [255, 136, 0], [255, 255, 0], [0, 255, 0], [0, 0, 255]
    ]
    body_connections = [
        ([1, 8], 9), ([1, 2], 8), ([1, 5], 7),
        ([2, 3], 10), ([3, 4], 11), ([5, 6], 6), ([6, 7], 5),
        ([8, 9], 12), ([9, 10], 13), ([1, 11], 4),
        ([11, 12], 3), ([12, 13], 2), ([1, 0], 1),
        ([0, 14], 16), ([14, 16], 18), ([0, 15], 15), ([15, 17], 17),
    ]

    def get_px(kp):
        return int(float(kp["x"]) * width), int(float(kp["y"]) * height)

    def is_vis(kp):
        return float(kp.get("vis", 0)) > 0.1

    stickwidth = 4
    for (i, j), color_idx in body_connections:
        if i < len(keypoints) and j < len(keypoints):
            kp_i, kp_j = keypoints[i], keypoints[j]
            if is_vis(kp_i) and is_vis(kp_j):
                p1, p2 = get_px(kp_i), get_px(kp_j)
                color = pose_colors[color_idx % len(pose_colors)]
                mX = np.mean([p1[0], p2[0]])
                mY = np.mean([p1[1], p2[1]])
                length = ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5
                angle = np.degrees(np.arctan2(p1[1] - p2[1], p1[0] - p2[0]))
                polygon = cv2.ellipse2Poly(
                    (int(mX), int(mY)), (int(length / 2), stickwidth), int(angle), 0, 360, 1
                )
                cv2.fillConvexPoly(canvas, polygon, color)

    for idx in range(min(18, len(keypoints))):
        kp = keypoints[idx]
        if is_vis(kp):
            cv2.circle(canvas, get_px(kp), 4, pose_colors[idx % len(pose_colors)], -1, cv2.LINE_AA)

    for body_wrist, hand_wrist in [(7, 88), (4, 109)]:
        if body_wrist < len(keypoints) and hand_wrist < len(keypoints):
            kp_b, kp_h = keypoints[body_wrist], keypoints[hand_wrist]
            if is_vis(kp_b) and is_vis(kp_h):
                cv2.line(canvas, get_px(kp_b), get_px(kp_h),
                         pose_colors[body_wrist % len(pose_colors)], 2, cv2.LINE_AA)

    for idx in range(FACE_START, min(FACE_START + FACE_COUNT, len(keypoints))):
        kp = keypoints[idx]
        if is_vis(kp):
            cv2.circle(canvas, get_px(kp), 2, [255, 255, 255], -1, cv2.LINE_AA)

    for hand_start in [LEFT_HAND_START, RIGHT_HAND_START]:
        finger_conns = []
        for fi in range(4):
            finger_conns.append((fi, fi + 1, 0))
        for fi in range(5, 8):
            finger_conns.append((fi, fi + 1, 1))
        for fi in range(9, 12):
            finger_conns.append((fi, fi + 1, 2))
        for fi in range(13, 16):
            finger_conns.append((fi, fi + 1, 3))
        for fi in range(17, 20):
            finger_conns.append((fi, fi + 1, 4))
        finger_conns += [(0, 5, 1), (0, 9, 2), (0, 13, 3), (0, 17, 4)]
        for (li, lj, fi) in finger_conns:
            gi, gj = hand_start + li, hand_start + lj
            if gi < len(keypoints) and gj < len(keypoints):
                kp_i, kp_j = keypoints[gi], keypoints[gj]
                if is_vis(kp_i) and is_vis(kp_j):
                    cv2.line(canvas, get_px(kp_i), get_px(kp_j), hand_finger_colors[fi], 2, cv2.LINE_AA)
        for hi in range(HAND_COUNT):
            gi = hand_start + hi
            if gi < len(keypoints) and is_vis(keypoints[gi]):
                if hi <= 4:
                    fi = 0
                elif hi <= 8:
                    fi = 1
                elif hi <= 12:
                    fi = 2
                elif hi <= 16:
                    fi = 3
                else:
                    fi = 4
                cv2.circle(canvas, get_px(keypoints[gi]), 3, hand_finger_colors[fi], -1, cv2.LINE_AA)

    return Image.fromarray(canvas)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=60)
def _run_pipeline(pose_image, prompt, lora_scale, num_inference_steps,
                  guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    """Run Krea-2 with a pose skeleton as the control/reference image.

    The generation resolution is matched to the pose image so the pose latents
    (placed on their own rotary grid from (0,0)) line up with the output.
    """
    generator = torch.Generator("cuda").manual_seed(int(seed))

    # Match the ComfyUI workflow's sizing EXACTLY (FluxKontextImageScale, node
    # 28): snap the pose skeleton to the nearest ~1 MP preferred bucket via
    # center-crop + Lanczos. The SAME scaled image is then used both as the VAE
    # reference (image=) and as the output canvas (height/width) — so the pose
    # reference-latent grid and the output grid share identical dimensions
    # (index_timestep_zero placement), giving strong pose adherence, and the
    # model runs at its native ~1 MP resolution for full quality.
    pose_image, out_w, out_h = _flux_kontext_scale(pose_image)

    # The pose LoRA (an ai-toolkit base-transformer LoRA, keys
    # diffusion_model.blocks.*.lora_{A,B}) is applied once at `lora_scale` via
    # the adapter weights — exactly like the workflow's LoraLoaderModelOnly
    # (strength_model=1.0). Do NOT also pass it through attention_kwargs["scale"]:
    # that double-scaled the LoRA and diverged from ComfyUI.
    pipe.transformer.set_adapters("default", weights=float(lora_scale))

    result = pipe(
        prompt=prompt,
        image=pose_image,
        height=out_h,
        width=out_w,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        generator=generator,
        # Root-cause fix for the blocky/checkerboard "patched noise" artifacts:
        # the pose skeleton's (mostly black) reference latents were bleeding into
        # the output as latent-grid-scale blocks. This LoRA was trained with
        # AI-Toolkit's `kv_cache` model kwarg (reference tokens attend ONLY to
        # each other), which is why the reference ComfyUI workflow
        # (krea2_controlnet_pose.json) sets Krea2OstrisEditModelPatch's kv_cache
        # widget to True. Running the default kv_cache=False path instead puts the
        # reference tokens in the per-step sequence with full bidirectional
        # attention — the wrong attention pattern for this LoRA — so the clean
        # reference latents leak into the generated tokens. Precomputing the
        # reference K/V once at t=0 and injecting them as isolated extra keys
        # (kv_cache=True) matches the trained/ComfyUI behaviour and removes the
        # artifacts.
        kv_cache=True,
    )
    return result.images[0]


def _resolve_seed(seed, randomize_seed):
    if randomize_seed:
        seed = int(np.random.randint(0, 2**31 - 1))
    return int(seed)


def generate_from_pose_image(
    pose_image,
    prompt,
    lora_scale=1.0,
    num_inference_steps=10,
    guidance_scale=0.0,
    seed=42,
    randomize_seed=True,
):
    """Generate using an already-rendered pose/skeleton image, used as-is.

    Args:
        pose_image: a pre-rendered OpenPose/DWPose skeleton image (on black).
        prompt: text describing the character, scene, and style to generate.
        lora_scale: strength of the pose-control LoRA (1.0 matches the workflow).
        num_inference_steps: denoising steps (10 matches the workflow's KSampler).
        guidance_scale: 0.0 = single forward, matching the workflow's cfg=1
            (no CFG) on this distilled Turbo model.
        seed: RNG seed for reproducibility.
        randomize_seed: pick a fresh seed each run.
    """
    if pose_image is None:
        raise gr.Error("Please provide a pose/skeleton image.")
    pose_image = (
        Image.fromarray(pose_image) if isinstance(pose_image, np.ndarray) else pose_image
    ).convert("RGB")
    seed = _resolve_seed(seed, randomize_seed)
    output = _run_pipeline(pose_image, prompt, lora_scale, num_inference_steps, guidance_scale, seed)
    return output, pose_image, seed


def generate_from_keypoints(
    keypoints,
    prompt,
    lora_scale=1.0,
    num_inference_steps=10,
    guidance_scale=0.0,
    seed=42,
    randomize_seed=True,
):
    """Generate using the (possibly user-edited) skeleton from the editor.

    Args:
        keypoints: normalized [{x, y, vis}, ...] skeleton from the editor.
        prompt: text describing the character, scene, and style to generate.
        lora_scale: strength of the pose-control LoRA (1.0 matches the workflow).
        num_inference_steps: denoising steps (10 matches the workflow's KSampler).
        guidance_scale: 0.0 = single forward, matching the workflow's cfg=1
            (no CFG) on this distilled Turbo model.
        seed: RNG seed for reproducibility.
        randomize_seed: pick a fresh seed each run.
    """
    if not keypoints or len(keypoints) == 0:
        raise gr.Error("No pose available. Upload a photo in the 'Regular Image' tab first.")

    # The editor renders on a 512x768 canvas; keep that aspect for the control
    # image so keypoints map cleanly to a portrait pose.
    pose_image = keypoints_to_pose_image(keypoints, 512, 768)
    seed = _resolve_seed(seed, randomize_seed)
    output = _run_pipeline(pose_image, prompt, lora_scale, num_inference_steps, guidance_scale, seed)
    return output, pose_image, seed


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLE_PROMPT = (
    "A high-resolution render of a futuristic cyber-ninja in a low-profile "
    "stealth pose. The suit is matte black with subtle blue and red LED "
    "accents. It holds two glowing plasma katanas low to the ground. The "
    "character is balanced on a high-rise ledge overlooking a rain-slicked "
    "cyberpunk megacity at night. Style: Cyberpunk aesthetics, realistic "
    "rendering, dramatic lighting."
)

with gr.Blocks(title="Krea 2 Pose ControlNet") as demo:
    gr.Markdown(
        """
        # 🎨 Krea 2 Turbo — Pose ControlNet LoRA
        Drive generation with a body pose. Either upload a ready-made pose
        skeleton, or upload a normal photo, auto-extract its skeleton, and drag
        the joints to tweak the pose before generating.
        Built on [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo)
        with the
        [thedeoxen/Krea-2-pose-controlnet](https://huggingface.co/thedeoxen/Krea-2-pose-controlnet)
        LoRA. Interactive editor inspired by
        [linoyts/Flux-2-control-pose](https://huggingface.co/spaces/linoyts/Flux-2-control-pose).
        """
    )

    # Shared prompt + advanced settings
    prompt = gr.Textbox(
        label="Prompt",
        placeholder="Describe the character, clothing, and scene…",
        value=EXAMPLE_PROMPT,
        lines=3,
    )
    with gr.Accordion("Advanced settings", open=False):
        lora_scale = gr.Slider(
            label="LoRA scale", minimum=0.0, maximum=1.5, step=0.05, value=1.0,
            info="Pose adherence strength. 1.0 matches the reference workflow.",
        )
        num_inference_steps = gr.Slider(
            label="Steps", minimum=4, maximum=20, step=1, value=10,
            info="Denoising steps. 10 matches the reference workflow's KSampler.",
        )
        guidance_scale = gr.Slider(
            label="Guidance scale", minimum=0.0, maximum=5.0, step=0.1, value=0.0,
            info="0 = single forward (matches the workflow's cfg=1 on Turbo).",
        )
        seed = gr.Number(label="Seed", value=42, precision=0)
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            with gr.Tabs():
                # --- Tab 1: pose skeleton used as-is ---
                with gr.Tab("Pose Image"):
                    gr.Markdown(
                        "Upload an **already-rendered pose/skeleton image**. "
                        "It is used as the control input exactly as provided."
                    )
                    pose_image_input = gr.Image(
                        label="Pose skeleton image", type="numpy", height=320,
                    )
                    run_pose_btn = gr.Button("Generate", variant="primary", size="lg")
                    gr.Examples(
                        examples=[
                            ["pose_example_1.png"],
                            ["pose_example_2.png"],
                        ],
                        inputs=[pose_image_input],
                    )

                # --- Tab 2: regular photo → auto skeleton → interactive edit ---
                with gr.Tab("Regular Image"):
                    gr.Markdown(
                        "Upload a **normal photo**. The skeleton is extracted "
                        "automatically — then **drag the joints** below to tweak "
                        "the pose before generating."
                    )
                    regular_image_input = gr.Image(
                        label="Photo (pose source)", type="pil", height=280,
                        sources=["upload", "webcam", "clipboard"],
                    )
                    pose_skeleton = InteractivePoseSkeleton(value=None, width=512, height=768)
                    pose_state = gr.State(None)
                    run_kp_btn = gr.Button("Generate", variant="primary", size="lg")

        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated image", type="pil", height=400)
            pose_preview = gr.Image(
                label="Pose used (what the model sees)", type="pil", height=240,
            )
            used_seed = gr.Number(label="Seed used", interactive=False)

    # --- Wiring: Pose Image tab ---
    run_pose_btn.click(
        fn=generate_from_pose_image,
        inputs=[pose_image_input, prompt, lora_scale, num_inference_steps,
                guidance_scale, seed, randomize_seed],
        outputs=[output_image, pose_preview, used_seed],
        api_name="generate_from_pose_image",
    )

    # --- Wiring: Regular Image tab ---
    def on_photo_upload(image):
        """Auto-extract a skeleton from an uploaded photo for the editor."""
        if image is None:
            return None, None
        gr.Info("🔍 Extracting pose from image…")
        keypoints, _ = extract_pose_keypoints(image)
        if keypoints is None:
            return None, None
        gr.Info("✅ Pose extracted — drag joints to edit, then Generate.")
        return keypoints, keypoints

    regular_image_input.change(
        fn=on_photo_upload,
        inputs=[regular_image_input],
        outputs=[pose_skeleton, pose_state],
    )

    # Keep the editor's edits in a State that survives the worker boundary.
    pose_skeleton.change(
        fn=lambda kp: kp,
        inputs=[pose_skeleton],
        outputs=[pose_state],
    )

    run_kp_btn.click(
        fn=generate_from_keypoints,
        inputs=[pose_state, prompt, lora_scale, num_inference_steps,
                guidance_scale, seed, randomize_seed],
        outputs=[output_image, pose_preview, used_seed],
        api_name="generate_from_keypoints",
    )


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)