multiview_runtime.py

"""Hosted Qwen studio generation through Hugging Face Inference Providers."""

from __future__ import annotations

import os
import time
import uuid
import zipfile
from pathlib import Path

import numpy as np
from PIL import Image


BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
ANGLE_MODEL_ID = "fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA"
OUTPUT_ROOT = Path("/tmp/studio10-outputs")

STUDIO_SHOTS = (
    ("01 · Front-left", None),
    ("02 · Front-right", "<sks> front-right quarter view elevated shot medium shot"),
    ("03 · Front", "<sks> front view eye-level shot medium shot"),
    ("04 · Left profile", "<sks> left side view eye-level shot medium shot"),
    ("05 · Right profile", "<sks> right side view eye-level shot medium shot"),
    ("06 · Rear-right", "<sks> back-right quarter view elevated shot medium shot"),
    ("07 · Rear-left", "<sks> back-left quarter view elevated shot medium shot"),
    ("08 · Low angle", "<sks> front-left quarter view low-angle shot medium shot"),
    ("09 · High angle", "<sks> front-right quarter view high-angle shot medium shot"),
    ("10 · Rear", "<sks> back view eye-level shot medium shot"),
)


def _client():
    from huggingface_hub import InferenceClient

    token = os.getenv("HF_TOKEN")
    if not token:
        raise RuntimeError("HF_TOKEN is not configured in the Space secrets.")
    return InferenceClient(provider="fal-ai", token=token, timeout=300)


def _identity_card(image: Image.Image, size: int = 1024) -> Image.Image:
    product = image.convert("RGBA")
    bbox = product.getchannel("A").getbbox()
    if bbox:
        product = product.crop(bbox)
    product.thumbnail((size - 180, size - 180), Image.Resampling.LANCZOS)
    card = Image.new("RGBA", (size, size), "white")
    x = (size - product.width) // 2
    y = size - product.height - 90
    card.alpha_composite(product, (x, y))
    return card.convert("RGB")


def _pure_white_finish(image: Image.Image) -> Image.Image:
    array = np.asarray(image.convert("RGB")).copy()
    low = array.min(axis=2)
    spread = array.max(axis=2) - low
    array[(low >= 247) & (spread <= 7)] = 255
    return Image.fromarray(array, mode="RGB")


def generate_studio_photos(
    isolated_reference: Image.Image,
    seed: int,
) -> list[tuple[str, Image.Image]]:
    client = _client()
    reference = _identity_card(isolated_reference)
    master_prompt = (
        "Create a photorealistic ecommerce studio photo of this exact product from a front-left "
        "three-quarter elevated camera angle. Preserve its precise shape, proportions, color, material, "
        "stitching, seams, hardware, logos and labels. Rebuild the whole photograph; do not paste the "
        "cutout. Place the product naturally on a seamless pure white studio floor with its real base "
        "fully touching the floor. Add softbox lighting and a short attached contact shadow. Never float "
        "or levitate the product. One product only, centered, fully visible, no props, no text, no border."
    )
    master = client.image_to_image(
        image=reference,
        prompt=master_prompt,
        model=BASE_MODEL_ID,
        num_inference_steps=40,
        guidance_scale=1.0,
        seed=int(seed),
    )
    master = _pure_white_finish(master)
    results: list[tuple[str, Image.Image]] = [(STUDIO_SHOTS[0][0], master)]

    for index, (title, pose_prompt) in enumerate(STUDIO_SHOTS[1:], start=1):
        output = client.image_to_image(
            image=master,
            prompt=pose_prompt,
            model=ANGLE_MODEL_ID,
            num_inference_steps=40,
            guidance_scale=1.0,
            seed=int(seed) + index * 997,
        )
        results.append((title, _pure_white_finish(output)))
    return results


def save_studio_outputs(results: list[tuple[str, Image.Image]]) -> tuple[list[str], str]:
    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    now = time.time()
    for directory in OUTPUT_ROOT.iterdir():
        try:
            if directory.is_dir() and now - directory.stat().st_mtime > 6 * 3600:
                for child in directory.iterdir():
                    child.unlink(missing_ok=True)
                directory.rmdir()
        except OSError:
            pass

    run_dir = OUTPUT_ROOT / uuid.uuid4().hex
    run_dir.mkdir()
    paths: list[str] = []
    for index, (_, image) in enumerate(results, start=1):
        path = run_dir / f"{index:02d}-studio-shot.png"
        image.save(path, optimize=True)
        paths.append(str(path))

    archive_path = run_dir / "studio10.zip"
    with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
        for path in paths:
            archive.write(path, arcname=Path(path).name)
    return paths, str(archive_path)