Fix open pose vector generation

Acly · Oct 15, 2024 · eaaa432 · eaaa432
1 parent 7775cc2
commit eaaa432
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/ai_diffusion/comfy_client.py b/ai_diffusion/comfy_client.py
@@ -727,7 +727,8 @@ def _extract_pose_json(msg: dict):
     try:
         output = msg["data"]["output"]
         if "openpose_json" in output:
-            return json.loads(output["openpose_json"][0])
+            result = json.loads(output["openpose_json"][0])
+            return result[0] if isinstance(result, list) else result
     except Exception as e:
         log.warning(f"Error processing message, error={str(e)}, msg={msg}")
     return None

diff --git a/tests/test_service.py b/tests/test_service.py
@@ -132,18 +132,13 @@ def test_large_image(qtapp, cloud_client):
     run_and_save(qtapp, cloud_client, workflow, "pod_large_image")
 
 
-@pytest.mark.parametrize("scenario", ["resolution", "steps", "control", "max_pixels"])
+@pytest.mark.parametrize("scenario", ["resolution", "steps", "max_pixels"])
 def test_validation(qtapp, cloud_client: CloudClient, scenario: str):
     workflow = create_simple_workflow()
     if scenario == "resolution":
         workflow.images = ImageInput.from_extent(Extent(19000, 512))
     elif scenario == "steps":
         ensure(workflow.sampling).total_steps = 200
-    elif scenario == "control":
-        img = Image.create(Extent(4, 4))
-        control = ensure(workflow.conditioning).control
-        for i in range(7):
-            control.append(ControlInput(ControlMode.depth, img))
     elif scenario == "max_pixels":
         workflow.images = ImageInput.from_extent(Extent(3840, 2168))  # > 4k