feat: decouple core photo cleaning logic for CLI and API reuse (#2)

cleaner.py

@@ -0,0 +1,222 @@
+"""Core garment photo cleaning logic for CLI and API reuse."""
+
+import base64
+from dataclasses import dataclass
+import io
+import mimetypes
+import os
+from typing import Any
+
+import requests
+
+API_URL = "https://openrouter.ai/api/v1/images"
+DEFAULT_MODEL = os.environ.get("MODEL", "meta/muse-image")
+
+BACKGROUND_PRESETS = {
+    "beige": "Soft warm beige studio background, evenly lit, no shadows, no props.",
+    "white": "Plain bright white studio backdrop, evenly lit, no shadows, no gradients, no props.",
+}
+
+EXTENSIONS = {
+    "image/png": ".png",
+    "image/jpeg": ".jpg",
+    "image/webp": ".webp",
+}
+
+
+def build_prompt(bg_choice: str) -> str:
+    bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice)
+    return (
+        "Product photography edit of this garment photo. Keep the garments pixel-faithful: identical "
+        "shape, colors, prints, labels, buttons, and every existing flaw unchanged. Remove ALL wrinkles "
+        "and creases completely: the fabric must look perfectly smooth and freshly ironed, flat like a "
+        f"new catalog product photo. {bg_description}"
+    )
+
+
+def encode_image_bytes(data: bytes, mime: str = "image/jpeg") -> str:
+    return f"data:{mime};base64,{base64.b64encode(data).decode()}"
+
+
+def encode_image(path: str) -> str:
+    mime = mimetypes.guess_type(path)[0] or "image/png"
+    with open(path, "rb") as f:
+        return encode_image_bytes(f.read(), mime=mime)
+
+
+def restore_resolution_bytes(
+    gen_bytes: bytes,
+    orig_width: int,
+    orig_height: int,
+    fmt: str = "JPEG",
+) -> bytes:
+    """Resize generated image bytes back to the original dimensions using Lanczos resampling."""
+    try:
+        from PIL import Image
+    except ImportError:
+        return gen_bytes
+
+    with Image.open(io.BytesIO(gen_bytes)) as gen:
+        if (gen.width, gen.height) == (orig_width, orig_height):
+            return gen_bytes
+        resample = getattr(Image, "Resampling", Image).LANCZOS
+        resized = gen.resize((orig_width, orig_height), resample)
+        buf = io.BytesIO()
+        kwargs = {"quality": 95} if fmt.upper() in ("JPEG", "JPG") else {}
+        resized.save(buf, format=fmt, **kwargs)
+        return buf.getvalue()
+
+
+def restore_resolution(out: str, src: str) -> None:
+    """Resize the generated image back to the input dimensions using Lanczos resampling."""
+    try:
+        from PIL import Image
+    except ImportError:
+        print("Pillow not installed; skipping resolution restore (pip install Pillow)")
+        return
+
+    with Image.open(src) as orig, Image.open(out) as gen:
+        if gen.size == orig.size:
+            return
+        resample = getattr(Image, "Resampling", Image).LANCZOS
+        resized = gen.resize(orig.size, resample)
+        ext = os.path.splitext(out)[1].lower()
+        fmt = "JPEG" if ext in (".jpg", ".jpeg") else None
+        kwargs = {"quality": 95} if fmt == "JPEG" else {}
+        if fmt:
+            resized.save(out, fmt, **kwargs)
+        else:
+            resized.save(out)
+        print(f"Restored resolution: {gen.size} -> {orig.size}")
+
+
+class CleanerError(Exception):
+    """Base error raised when cleaning or calling OpenRouter fails."""
+    pass
+
+
+@dataclass
+class CleanResult:
+    image_bytes: bytes
+    media_type: str = "image/jpeg"
+    cost: float | None = None
+    width: int | None = None
+    height: int | None = None
+    original_dimensions: tuple[int, int] | None = None
+    model_dimensions: tuple[int, int] | None = None
+
+
+def call_openrouter_images(
+    image_url_or_data_uri: str,
+    prompt: str,
+    api_key: str,
+    model: str = DEFAULT_MODEL,
+    api_url: str = API_URL,
+    timeout: int = 300,
+) -> dict:
+    """Submit image edit request to OpenRouter images API."""
+    headers = {"Authorization": f"Bearer {api_key}"}
+    payload: dict[str, Any] = {
+        "model": model,
+        "prompt": prompt,
+        "background": "opaque",
+        "input_references": [
+            {"type": "image_url", "image_url": {"url": image_url_or_data_uri}}
+        ],
+        "output_format": "jpeg",
+    }
+    resp = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
+    if resp.status_code != 200:
+        try:
+            message = resp.json()["error"]["message"]
+        except Exception:
+            message = resp.text
+        raise CleanerError(f"API error {resp.status_code}: {message}")
+    return resp.json()
+
+
+def clean_garment(
+    image_data: bytes | str,
+    prompt: str,
+    api_key: str,
+    model: str = DEFAULT_MODEL,
+    restore_res: bool = True,
+    mime: str = "image/jpeg",
+    api_url: str = API_URL,
+    timeout: int = 300,
+) -> CleanResult:
+    """Clean garment image given file path or bytes, returning CleanResult."""
+    orig_width: int | None = None
+    orig_height: int | None = None
+
+    if isinstance(image_data, bytes):
+        try:
+            from PIL import Image
+            with Image.open(io.BytesIO(image_data)) as im:
+                orig_width, orig_height = im.size
+                if im.format:
+                    mime = Image.MIME.get(im.format, mime)
+        except Exception:
+            pass
+        data_uri = encode_image_bytes(image_data, mime=mime)
+    elif isinstance(image_data, str):
+        if image_data.startswith("data:"):
+            data_uri = image_data
+        else:
+            data_uri = encode_image(image_data)
+            try:
+                from PIL import Image
+                with Image.open(image_data) as im:
+                    orig_width, orig_height = im.size
+            except Exception:
+                pass
+    else:
+        raise TypeError("image_data must be bytes or str path")
+
+    result = call_openrouter_images(
+        image_url_or_data_uri=data_uri,
+        prompt=prompt,
+        api_key=api_key,
+        model=model,
+        api_url=api_url,
+        timeout=timeout,
+    )
+
+    image_entry = result["data"][0]
+    gen_bytes = base64.b64decode(image_entry["b64_json"])
+    media_type = image_entry.get("media_type", "image/jpeg")
+    cost = result.get("usage", {}).get("cost")
+    model_dimensions: tuple[int, int] | None = None
+    try:
+        from PIL import Image
+        with Image.open(io.BytesIO(gen_bytes)) as im:
+            model_dimensions = im.size
+    except Exception:
+        pass
+
+    if restore_res and orig_width is not None and orig_height is not None:
+        fmt = "JPEG" if media_type == "image/jpeg" else "PNG"
+        gen_bytes = restore_resolution_bytes(
+            gen_bytes, orig_width, orig_height, fmt=fmt
+        )
+
+    width: int | None = orig_width if (restore_res and orig_width is not None) else (model_dimensions[0] if model_dimensions else None)
+    height: int | None = orig_height if (restore_res and orig_height is not None) else (model_dimensions[1] if model_dimensions else None)
+    try:
+        from PIL import Image
+        with Image.open(io.BytesIO(gen_bytes)) as im:
+            width, height = im.size
+    except Exception:
+        pass
+
+    orig_dimensions = (orig_width, orig_height) if (orig_width is not None and orig_height is not None) else None
+
+    return CleanResult(
+        image_bytes=gen_bytes,
+        media_type=media_type,
+        cost=cost,
+        width=width,
+        height=height,
+        original_dimensions=orig_dimensions,
+        model_dimensions=model_dimensions,
+    )

tests/test_cleaner.py

@@ -0,0 +1,287 @@
+import pytest
+from cleaner import (
+    API_URL,
+    DEFAULT_MODEL,
+    BACKGROUND_PRESETS,
+    EXTENSIONS,
+    build_prompt,
+    encode_image,
+    encode_image_bytes,
+    restore_resolution,
+    restore_resolution_bytes,
+    call_openrouter_images,
+    clean_garment,
+    CleanResult,
+    CleanerError,
+)
+import io
+from PIL import Image
+
+
+def test_constants():
+    assert API_URL == "https://openrouter.ai/api/v1/images"
+    assert "beige" in BACKGROUND_PRESETS
+    assert "white" in BACKGROUND_PRESETS
+    assert EXTENSIONS["image/jpeg"] == ".jpg"
+    assert EXTENSIONS["image/png"] == ".png"
+    assert EXTENSIONS["image/webp"] == ".webp"
+
+
+def test_build_prompt_presets():
+    prompt_beige = build_prompt("beige")
+    assert BACKGROUND_PRESETS["beige"] in prompt_beige
+    assert "Remove ALL wrinkles" in prompt_beige
+
+    prompt_white = build_prompt("white")
+    assert BACKGROUND_PRESETS["white"] in prompt_white
+    assert "Remove ALL wrinkles" in prompt_white
+
+    # Case insensitivity
+    assert build_prompt("BEIGE") == prompt_beige
+
+
+def test_build_prompt_custom():
+    custom = "Vintage wooden floor with warm sidelight."
+    prompt_custom = build_prompt(custom)
+    assert custom in prompt_custom
+    assert "Remove ALL wrinkles" in prompt_custom
+
+
+def test_encode_image_bytes():
+    raw = b"fake-image-bytes"
+    encoded = encode_image_bytes(raw, "image/jpeg")
+    assert encoded == "data:image/jpeg;base64,ZmFrZS1pbWFnZS1ieXRlcw=="
+
+    # Default mime
+    encoded_default = encode_image_bytes(raw)
+    assert encoded_default.startswith("data:image/jpeg;base64,")
+
+
+def test_encode_image_file(tmp_path):
+    img_file = tmp_path / "test.png"
+    img_file.write_bytes(b"png-data")
+    encoded = encode_image(str(img_file))
+    assert encoded == "data:image/png;base64,cG5nLWRhdGE="
+
+    jpg_file = tmp_path / "test.jpg"
+    jpg_file.write_bytes(b"jpg-data")
+    encoded_jpg = encode_image(str(jpg_file))
+    assert encoded_jpg == "data:image/jpeg;base64,anBnLWRhdGE="
+
+    unknown_file = tmp_path / "test.customext"
+    unknown_file.write_bytes(b"custom-data")
+    encoded_unknown = encode_image(str(unknown_file))
+    assert encoded_unknown.startswith("data:image/png;base64,")
+
+
+def create_test_image(size=(100, 100), fmt="JPEG"):
+    buf = io.BytesIO()
+    img = Image.new("RGB", size, color="blue")
+    img.save(buf, format=fmt)
+    return buf.getvalue()
+
+
+def test_restore_resolution_bytes_different_dimensions():
+    gen_data = create_test_image(size=(100, 100), fmt="JPEG")
+    orig_width, orig_height = 200, 300
+
+    restored = restore_resolution_bytes(gen_data, orig_width, orig_height, fmt="JPEG")
+    assert restored != gen_data
+
+    with Image.open(io.BytesIO(restored)) as im:
+        assert im.size == (200, 300)
+
+
+def test_restore_resolution_bytes_same_dimensions_noop():
+    gen_data = create_test_image(size=(100, 100), fmt="JPEG")
+    restored = restore_resolution_bytes(gen_data, 100, 100, fmt="JPEG")
+    assert restored is gen_data  # Exact same object/bytes, no re-encoding
+
+
+def test_restore_resolution_files(tmp_path):
+    orig_file = tmp_path / "orig.jpg"
+    out_file = tmp_path / "out.jpg"
+
+    orig_img = Image.new("RGB", (300, 400), color="red")
+    orig_img.save(str(orig_file), format="JPEG")
+
+    gen_img = Image.new("RGB", (150, 200), color="green")
+    gen_img.save(str(out_file), format="JPEG")
+
+    restore_resolution(str(out_file), str(orig_file))
+
+    with Image.open(str(out_file)) as im:
+        assert im.size == (300, 400)
+
+
+def test_restore_resolution_files_same_dimensions(tmp_path):
+    orig_file = tmp_path / "orig.jpg"
+    out_file = tmp_path / "out.jpg"
+
+    orig_img = Image.new("RGB", (200, 200), color="red")
+    orig_img.save(str(orig_file), format="JPEG")
+    orig_img.save(str(out_file), format="JPEG")
+
+    mtime_before = out_file.stat().st_mtime_ns
+    restore_resolution(str(out_file), str(orig_file))
+    mtime_after = out_file.stat().st_mtime_ns
+
+    # File should not have been overwritten
+    assert mtime_before == mtime_after
+    with Image.open(str(out_file)) as im:
+        assert im.size == (200, 200)
+
+
+def test_call_openrouter_images_success(monkeypatch):
+    def mock_post(url, headers, json, timeout):
+        assert url == API_URL
+        assert headers == {"Authorization": "Bearer test-key"}
+        assert json["model"] == DEFAULT_MODEL
+        assert json["prompt"] == "test prompt"
+        assert json["background"] == "opaque"
+        assert json["output_format"] == "jpeg"
+        assert json["input_references"][0]["image_url"]["url"] == "data:image/jpeg;base64,abc"
+        assert timeout == 300
+
+        class MockResponse:
+            status_code = 200
+
+            def json(self):
+                return {
+                    "data": [{"b64_json": "ZGF0YQ==", "media_type": "image/jpeg"}],
+                    "usage": {"cost": 0.05},
+                }
+
+        return MockResponse()
+
+    import requests
+    monkeypatch.setattr(requests, "post", mock_post)
+
+    result = call_openrouter_images(
+        image_url_or_data_uri="data:image/jpeg;base64,abc",
+        prompt="test prompt",
+        api_key="test-key",
+    )
+    assert result["data"][0]["b64_json"] == "ZGF0YQ=="
+    assert result["usage"]["cost"] == 0.05
+
+
+def test_call_openrouter_images_error_json(monkeypatch):
+    def mock_post(url, headers, json, timeout):
+        class MockResponse:
+            status_code = 401
+            text = '{"error": {"message": "Invalid API key"}}'
+
+            def json(self):
+                return {"error": {"message": "Invalid API key"}}
+
+        return MockResponse()
+
+    import requests
+    monkeypatch.setattr(requests, "post", mock_post)
+
+    with pytest.raises(CleanerError) as exc_info:
+        call_openrouter_images(
+            image_url_or_data_uri="data:image/jpeg;base64,abc",
+            prompt="test prompt",
+            api_key="invalid-key",
+        )
+    assert "API error 401: Invalid API key" in str(exc_info.value)
+
+
+def test_call_openrouter_images_error_text(monkeypatch):
+    def mock_post(url, headers, json, timeout):
+        class MockResponse:
+            status_code = 502
+            text = "Bad Gateway"
+
+            def json(self):
+                raise ValueError("Not JSON")
+
+        return MockResponse()
+
+    import requests
+    monkeypatch.setattr(requests, "post", mock_post)
+
+    with pytest.raises(CleanerError) as exc_info:
+        call_openrouter_images(
+            image_url_or_data_uri="data:image/jpeg;base64,abc",
+            prompt="test prompt",
+            api_key="test-key",
+        )
+    assert "API error 502: Bad Gateway" in str(exc_info.value)
+
+
+def test_clean_garment_bytes(monkeypatch):
+    orig_bytes = create_test_image(size=(300, 400), fmt="JPEG")
+    generated_bytes = create_test_image(size=(150, 200), fmt="JPEG")
+    import base64
+    b64_gen = base64.b64encode(generated_bytes).decode()
+
+    def mock_call_openrouter(image_url_or_data_uri, prompt, api_key, model=DEFAULT_MODEL, api_url=API_URL, timeout=300):
+        assert image_url_or_data_uri.startswith("data:image/jpeg;base64,")
+        assert prompt == "test prompt"
+        assert api_key == "test-key"
+        return {
+            "data": [{"b64_json": b64_gen, "media_type": "image/jpeg"}],
+            "usage": {"cost": 0.02},
+        }
+
+    import cleaner
+    monkeypatch.setattr(cleaner, "call_openrouter_images", mock_call_openrouter)
+
+    # With restore_res=True
+    res = clean_garment(
+        image_data=orig_bytes,
+        prompt="test prompt",
+        api_key="test-key",
+        restore_res=True,
+    )
+    assert isinstance(res, CleanResult)
+    assert res.media_type == "image/jpeg"
+    assert res.cost == 0.02
+    assert res.width == 300
+    assert res.height == 400
+    with Image.open(io.BytesIO(res.image_bytes)) as im:
+        assert im.size == (300, 400)
+
+    # With restore_res=False
+    res_unscaled = clean_garment(
+        image_data=orig_bytes,
+        prompt="test prompt",
+        api_key="test-key",
+        restore_res=False,
+    )
+    assert res_unscaled.width == 150
+    assert res_unscaled.height == 200
+    with Image.open(io.BytesIO(res_unscaled.image_bytes)) as im:
+        assert im.size == (150, 200)
+
+
+def test_clean_garment_file(tmp_path, monkeypatch):
+    orig_file = tmp_path / "garment.jpg"
+    orig_img = Image.new("RGB", (250, 350), color="yellow")
+    orig_img.save(str(orig_file), format="JPEG")
+
+    gen_bytes = create_test_image(size=(100, 100), fmt="JPEG")
+    import base64
+    b64_gen = base64.b64encode(gen_bytes).decode()
+
+    def mock_call_openrouter(image_url_or_data_uri, prompt, api_key, model=DEFAULT_MODEL, api_url=API_URL, timeout=300):
+        return {
+            "data": [{"b64_json": b64_gen, "media_type": "image/jpeg"}],
+            "usage": {"cost": 0.03},
+        }
+
+    import cleaner
+    monkeypatch.setattr(cleaner, "call_openrouter_images", mock_call_openrouter)
+
+    res = clean_garment(
+        image_data=str(orig_file),
+        prompt="test prompt",
+        api_key="test-key",
+        restore_res=True,
+    )
+    assert res.width == 250
+    assert res.height == 350
+    assert res.cost == 0.03

tests/test_unwrap_clothes.py

@@ -0,0 +1,130 @@
+import io
+import subprocess
+import sys
+
+from PIL import Image
+import pytest
+
+import unwrap_clothes
+
+
+def test_reexported_symbols():
+    # Ensure backwards compatibility for any caller importing from unwrap_clothes
+    assert hasattr(unwrap_clothes, "API_URL")
+    assert hasattr(unwrap_clothes, "DEFAULT_MODEL")
+    assert hasattr(unwrap_clothes, "BACKGROUND_PRESETS")
+    assert hasattr(unwrap_clothes, "EXTENSIONS")
+    assert hasattr(unwrap_clothes, "build_prompt")
+    assert hasattr(unwrap_clothes, "encode_image")
+    assert hasattr(unwrap_clothes, "restore_resolution")
+
+
+def test_cli_help():
+    result = subprocess.run(
+        [sys.executable, "unwrap_clothes.py", "--help"],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0
+    assert "Unwrinkle clothing photos and unify backgrounds via OpenRouter." in result.stdout
+    assert "--bg" in result.stdout
+    assert "--model" in result.stdout
+    assert "--api-key" in result.stdout
+    assert "--no-restore-res" in result.stdout
+
+
+def test_main_file_not_found(monkeypatch, capsys):
+    monkeypatch.setattr(sys, "argv", ["unwrap_clothes.py", "non_existent_file.jpg", "--api-key", "key"])
+    with pytest.raises(SystemExit) as exc:
+        unwrap_clothes.main()
+    assert "Error: file not found: non_existent_file.jpg" in str(exc.value)
+
+
+def test_main_missing_api_key(tmp_path, monkeypatch, capsys):
+    img = tmp_path / "img.jpg"
+    img.write_bytes(b"data")
+    monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
+    monkeypatch.setattr(sys, "argv", ["unwrap_clothes.py", str(img)])
+    with pytest.raises(SystemExit) as exc:
+        unwrap_clothes.main()
+    assert "missing API key" in str(exc.value)
+
+
+def test_main_success_flow(tmp_path, monkeypatch, capsys):
+    # Setup test image
+    src_file = tmp_path / "source.jpg"
+    im = Image.new("RGB", (200, 300), color="blue")
+    im.save(str(src_file), format="JPEG")
+
+    out_file = tmp_path / "custom_out.jpg"
+
+    from cleaner import CleanResult
+
+    cleaned_img = Image.new("RGB", (200, 300), color="white")
+    buf = io.BytesIO()
+    cleaned_img.save(buf, format="JPEG")
+    fake_result = CleanResult(
+        image_bytes=buf.getvalue(),
+        media_type="image/jpeg",
+        cost=0.04,
+        width=200,
+        height=300,
+    )
+
+    clean_garment_called = {}
+
+    def mock_clean_garment(image_data, prompt, api_key, model, restore_res):
+        clean_garment_called["image_data"] = image_data
+        clean_garment_called["prompt"] = prompt
+        clean_garment_called["api_key"] = api_key
+        clean_garment_called["model"] = model
+        clean_garment_called["restore_res"] = restore_res
+        return fake_result
+
+    monkeypatch.setattr("unwrap_clothes.clean_garment", mock_clean_garment)
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "unwrap_clothes.py",
+            str(src_file),
+            "-o",
+            str(out_file),
+            "--bg",
+            "beige",
+            "--api-key",
+            "test-token",
+        ],
+    )
+
+    unwrap_clothes.main()
+
+    assert clean_garment_called["api_key"] == "test-token"
+    assert clean_garment_called["restore_res"] is True
+    assert out_file.exists()
+    assert out_file.read_bytes() == fake_result.image_bytes
+
+    captured = capsys.readouterr()
+    assert f"Submitting to {unwrap_clothes.DEFAULT_MODEL}..." in captured.out
+    assert f"Saved: {out_file} (cost $0.04)" in captured.out
+
+
+def test_main_api_error_handling(tmp_path, monkeypatch):
+    src_file = tmp_path / "source.jpg"
+    src_file.write_bytes(b"dummy")
+
+    from cleaner import CleanerError
+
+    def mock_clean_garment(*args, **kwargs):
+        raise CleanerError("API error 400: Invalid image format")
+
+    monkeypatch.setattr("unwrap_clothes.clean_garment", mock_clean_garment)
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        ["unwrap_clothes.py", str(src_file), "--api-key", "token"],
+    )
+
+    with pytest.raises(SystemExit) as exc:
+        unwrap_clothes.main()
+    assert "API error 400: Invalid image format" in str(exc.value)

unwrap_clothes.py

@@ -6,64 +6,40 @@ and replacing the background with a studio backdrop.
 """
 
 import argparse
-import base64
-import mimetypes
 import os
 import sys
 
-import requests
+from cleaner import (
+    API_URL,
+    BACKGROUND_PRESETS,
+    CleanerError,
+    CleanResult,
+    DEFAULT_MODEL,
+    EXTENSIONS,
+    build_prompt,
+    clean_garment,
+    encode_image,
+    encode_image_bytes,
+    restore_resolution,
+    restore_resolution_bytes,
+)
 
-API_URL = "https://openrouter.ai/api/v1/images"
+__all__ = [
-DEFAULT_MODEL = os.environ.get("MODEL", "meta/muse-image")
+    "API_URL",
-
+    "DEFAULT_MODEL",
-BACKGROUND_PRESETS = {
+    "BACKGROUND_PRESETS",
-    "beige": "Soft warm beige studio background, evenly lit, no shadows, no props.",
+    "EXTENSIONS",
-    "white": "Plain bright white studio backdrop, evenly lit, no shadows, no gradients, no props.",
+    "build_prompt",
-}
+    "encode_image",
-
+    "encode_image_bytes",
-EXTENSIONS = {
+    "restore_resolution",
-    "image/png": ".png",
+    "restore_resolution_bytes",
-    "image/jpeg": ".jpg",
+    "clean_garment",
-    "image/webp": ".webp",
+    "CleanResult",
-}
+    "CleanerError",
-
+    "parse_args",
-
+    "main",
-def build_prompt(bg_choice: str) -> str:
+]
-    bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice)
-    return (
-        "Product photography edit of this garment photo. Keep the garments pixel-faithful: identical "
-        "shape, colors, prints, labels, buttons, and every existing flaw unchanged. Remove ALL wrinkles "
-        "and creases completely: the fabric must look perfectly smooth and freshly ironed, flat like a "
-        f"new catalog product photo. {bg_description}"
-    )
-
-
-def encode_image(path: str) -> str:
-    mime = mimetypes.guess_type(path)[0] or "image/png"
-    with open(path, "rb") as f:
-        return f"data:{mime};base64,{base64.b64encode(f.read()).decode()}"
-
-
-def restore_resolution(out: str, src: str) -> None:
-    """Resize the generated image back to the input dimensions using Lanczos resampling."""
-    try:
-        from PIL import Image
-    except ImportError:
-        print("Pillow not installed; skipping resolution restore (pip install Pillow)")
-        return
-
-    with Image.open(src) as orig, Image.open(out) as gen:
-        if gen.size == orig.size:
-            return
-        resized = gen.resize(orig.size, Image.LANCZOS)
-        ext = os.path.splitext(out)[1].lower()
-        fmt = "JPEG" if ext in (".jpg", ".jpeg") else None
-        kwargs = {"quality": 95} if fmt == "JPEG" else {}
-        if fmt:
-            resized.save(out, fmt, **kwargs)
-        else:
-            resized.save(out)
-        print(f"Restored resolution: {gen.size} -> {orig.size}")
 
 
 def parse_args() -> argparse.Namespace:
@@ -113,44 +89,35 @@ def main() -> None:
     prompt = build_prompt(args.bg)
 
     print(f"Submitting to {args.model}...")
-    resp = requests.post(
+    try:
-        API_URL,
+        result = clean_garment(
-        headers={"Authorization": f"Bearer {args.api_key}"},
+            image_data=args.image,
-        json={
+            prompt=prompt,
-            "model": args.model,
+            api_key=args.api_key,
-            "prompt": prompt,
+            model=args.model,
-            "background": "opaque",
+            restore_res=not args.no_restore_res,
-            "input_references": [
+        )
-                {"type": "image_url", "image_url": {"url": encode_image(args.image)}}
+    except CleanerError as e:
-            ],
+        sys.exit(str(e))
-            "output_format": "jpeg",
-        },
-        timeout=300,
-    )
-
-    if resp.status_code != 200:
-        try:
-            message = resp.json()["error"]["message"]
-        except Exception:
-            message = resp.text
-        sys.exit(f"API error {resp.status_code}: {message}")
-
-    result = resp.json()
-    image = result["data"][0]
-    ext = EXTENSIONS.get(image.get("media_type"), ".jpg")
 
+    ext = EXTENSIONS.get(result.media_type, ".jpg")
     out = args.output or (os.path.splitext(args.image)[0] + "_clean" + ext)
 
     with open(out, "wb") as f:
-        f.write(base64.b64decode(image["b64_json"]))
+        f.write(result.image_bytes)
 
-    if not args.no_restore_res:
+    if (
-        restore_resolution(out, args.image)
+        not args.no_restore_res
+        and result.model_dimensions is not None
+        and result.original_dimensions is not None
+        and result.model_dimensions != result.original_dimensions
+    ):
+        print(f"Restored resolution: {result.model_dimensions} -> {result.original_dimensions}")
 
-    cost = result.get("usage", {}).get("cost")
+    cost = result.cost
     cost_str = f" (cost ${cost:.2f})" if cost is not None else ""
     print(f"Saved: {out}{cost_str}")
 
 
 if __name__ == "__main__":
-    main()
+    main()