"""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 = """
""" 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)