Add Pose lock mode (ControlNet OpenPose) for SD1.5 (#13)

- Add Pose lock mode (ControlNet OpenPose) for SD1.5 (337b47f6340c5259fe86f963b35faf621d30f193)


Co-authored-by: pormungtailaw <pormungtai@users.noreply.huggingface.co>

pipeline_manager.py

@@ -38,7 +38,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Modes supported per base family. Used by the UI to gate options.
 SUPPORTED_MODES = {
-    "sd15": ["txt2img", "img2img", "ip_adapter", "face_id"],
+    "sd15": ["txt2img", "img2img", "ip_adapter", "face_id", "pose"],
     "sdxl": ["txt2img", "img2img", "ip_adapter", "face_id"],
     "flux": ["txt2img", "img2img"],
 }
@@ -48,6 +48,7 @@ MODE_LABELS = {
     "img2img": "Image → Image (denoise)",
     "ip_adapter": "IP-Adapter (style / subject)",
     "face_id": "Face identity (FaceID)",
+    "pose": "Pose lock (ControlNet OpenPose)",
 }
 
 # ---------------------------------------------------------------------------
@@ -328,6 +329,34 @@ def _face_embeds(image):
 # ---------------------------------------------------------------------------
 # Generation
 # ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+# ControlNet (OpenPose) — locks the generated subject to an uploaded pose.
+# ---------------------------------------------------------------------------
+_CONTROLNET = {}
+_OPENPOSE = None
+
+
+def _get_controlnet(base):
+    if base in _CONTROLNET:
+        return _CONTROLNET[base]
+    from diffusers import ControlNetModel
+    if base == "sd15":
+        repo = "lllyasviel/control_v11p_sd15_openpose"
+    else:
+        raise ValueError("Pose (ControlNet) รองรับเฉพาะ SD1.5 ตอนนี้ / SD1.5 only for now.")
+    cn = ControlNetModel.from_pretrained(repo, torch_dtype=DTYPE_SD)
+    _CONTROLNET[base] = cn
+    return cn
+
+
+def _get_openpose():
+    global _OPENPOSE
+    if _OPENPOSE is None:
+        from controlnet_aux import OpenposeDetector
+        _OPENPOSE = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+    return _OPENPOSE
+
+
 def _safe_call(pipe_obj, call):
     """Run the pipeline; if clip_skip trips a version incompatibility, retry without it."""
     try:
@@ -411,4 +440,16 @@ def run_generation(cfg, mode, prompt, negative_prompt, ref_image,
         embeds = _face_embeds(ref_image).to(DEVICE)
         call["ip_adapter_image_embeds"] = [embeds]
 
+    elif mode == "pose":
+        if ref_image is None:
+            raise ValueError("Pose ต้องอัปโหลดรูปท่าทางก่อน / Upload a pose reference image first.")
+        _ensure_adapter(pipe, base, None)
+        detector = _get_openpose()
+        pose_img = detector(ref_image.convert("RGB")).resize((int(width), int(height)))
+        cn = _get_controlnet(base).to(DEVICE)
+        from diffusers import StableDiffusionControlNetPipeline
+        cn_pipe = StableDiffusionControlNetPipeline.from_pipe(pipe, controlnet=cn).to(DEVICE)
+        call["image"] = pose_img
+        return _safe_call(cn_pipe, call)
+
     return _safe_call(pipe, call)

requirements.txt

@@ -19,3 +19,5 @@ opencv-python-headless
 # Face-identity mode (IP-Adapter FaceID). Heavy; comment out if you don't use Face mode.
 insightface==0.7.3
 onnxruntime
+# Pose lock mode (ControlNet OpenPose) — pose detection from the uploaded image.
+controlnet_aux