Fix pose fidelity/quality: replicate ComfyUI FluxKontextImageScale (nearest ~1MP bucket + center-crop + Lanczos) for pose sizing so ref-latent grid == output grid at native ~1MP resolution

app.py

@@ -85,23 +85,59 @@ def _snap16(v):
     return max(512, min(1024, v))
 
 
-# Match the ComfyUI reference workflow's sizing: the pose image is scaled to
-# fit ~1 MP (FluxKontextImageScale), its size drives EmptyLatentImage
-# (GetImageSize), and the Krea-2 latent grid patchifies on a multiple of 16
-# (vae_scale_factor 8 * patch 2). Preserve aspect, fit under 1 MP, snap /16 —
-# NOT a hard 512/1024 clamp, which distorted the aspect and hurt pose fidelity.
-REF_MAX_PIXELS = 1024 * 1024
-REF_SNAP = 16
-
+# 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
+    )
 
-def _fit_pose_size(w, h):
-    """Scale (w, h) down to fit REF_MAX_PIXELS, keep aspect, snap to /16."""
-    import math
+    # 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))
 
-    scale = min(1.0, math.sqrt(REF_MAX_PIXELS / float(w * h)))
-    nw = max(int(round(w * scale / REF_SNAP)) * REF_SNAP, REF_SNAP)
-    nh = max(int(round(h * scale / REF_SNAP)) * REF_SNAP, REF_SNAP)
-    return nw, nh
+    # 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):
@@ -871,11 +907,14 @@ def _run_pipeline(pose_image, prompt, lora_scale, num_inference_steps,
     """
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    # Match the ComfyUI workflow's sizing: fit the pose image under ~1 MP,
-    # preserve aspect, snap to /16. The output grid then lines up with the pose
-    # reference latent grid (index_timestep_zero placement) — strong adherence.
-    w, h = pose_image.size
-    out_w, out_h = _fit_pose_size(w, h)
+    # 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