studio/refine.py

"""Astra visual inspection and run-isolated fx repair planning."""
import base64
import json
import os
from pathlib import Path
import re
import subprocess
import shutil

import httpx
from PIL import Image

from .contracts import Refinement
from .mesh import apply_refinement, describe, export
from .render import render

MODEL = "openai/gpt-6-astra"
FX_RUNTIME = Path(__file__).parent / "fx_runtime"
ROOT = Path(__file__).resolve().parents[1]
RUNTIME = ROOT / ".runtime"


def preflight():
    if not os.environ.get("OPENROUTER_API_KEY"):
        raise RuntimeError("Set the OPENROUTER_API_KEY Space secret before generating.")
    if not shutil.which("node") or not (FX_RUNTIME / "node_modules/libfx/package.json").is_file():
        raise RuntimeError("Install Node 20+ and run npm ci --prefix studio/fx_runtime --ignore-scripts.")
    from .render import check_renderer
    check_renderer()


def image_data_url(path):
    path = Path(path)
    with Image.open(path) as image:
        mime = Image.MIME.get(image.format)
    if mime not in {"image/png", "image/jpeg", "image/webp"}:
        raise ValueError("Astra evidence must be PNG, JPEG, or WebP.")
    return f"data:{mime};base64," + base64.b64encode(path.read_bytes()).decode()


def vision(prompt, images):
    content = [{"type": "text", "text": prompt}]
    for path in images:
        content.append({"type": "image_url", "image_url": {"url": image_data_url(path)}})
    with httpx.Client(timeout=120) as client:
        response = client.post("https://openrouter.ai/api/v1/chat/completions",
                               headers={"Authorization": "Bearer " + os.environ["OPENROUTER_API_KEY"]},
                               json={"model": MODEL, "messages": [{"role": "user", "content": content}],
                                     "max_tokens": 8192, "reasoning": {"effort": "low"}, "temperature": 0.2})
        if response.status_code != 200:
            raise RuntimeError(f"Astra vision endpoint returned HTTP {response.status_code}; no substitute model used.")
        payload = response.json()
        if payload.get("model") != MODEL:
            raise RuntimeError("The visual reviewer did not return the requested Astra model.")
        result = payload["choices"][0]["message"]["content"]
        if not isinstance(result, str) or not result.strip():
            raise RuntimeError("Astra returned an empty visual assessment; no refinement was accepted.")
        return result


def json_object(text):
    text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip())
    return json.loads(text)


def plan_with_fx(brief, parts, assessment, folder, images=(), previous=None):
    skill = RUNTIME / "orc-skills/game-dev/game-model-cleanup"
    instructions = (skill / "SKILL.md").read_text() + "\n" + (skill / "WORKFLOW.md").read_text()
    prompt = f"""Apply the game-model-cleanup skill below to this asset. This application's contract replaces the
skill's Final Stand-specific paths, roster, animation and equipment requirements. The source is immutable.
Inspect /work/parts.json, /work/assessment.txt, /work/schema.json and the cleanup skill using the virtual
bash tool. Reference and before renders are attached. Return only a JSON repair plan matching the schema.
Do not invent missing anatomy or features. All commands execute only in this run's in-memory filesystem.
You may rename existing parts, remove duplicate/degenerate faces, propose small PBR/color corrections,
and set project_reference_texture=true for evidence-supported reference texture transfer. This operator
projects only visible foreground reference pixels, retains occluded/back-facing texels and all PBR channels,
and never changes geometry or UVs. Use it for visible front texture mismatch when reference pose and model
silhouette align. Do not use projection for different poses or to repair geometry. reference_front_axis is
the outward viewing direction AFTER Blender glTF import, with Z up; canonical front is -Y. It is NOT the
source glTF Y-up frame. Choose +Y/-X/+X only if the rendered orientation establishes that front.
Use projection_strength between 0 and 1; 0.9 is appropriate for strong matching-reference correction.
Color gains are multiplicative RGB gains, each between 0.75 and 1.25 inclusive, not target RGB colors.
Use [1, 1, 1] to preserve the original color.
Preserve UVs, normals, geometry identity, scale and all components. Put unsupported geometry reconstruction,
unseen-surface repainting or rigging needs in unresolved. Do not mark supported projection as unsupported. Empty edits are appropriate when evidence supports no change.
Brief: {brief}
Parts: {json.dumps(describe(parts))}
Astra visual inspection: {assessment}
Schema: {json.dumps(Refinement.model_json_schema())}
Skill and workflow:\n{instructions}"""
    env = {key: value for key, value in os.environ.items() if key in {
        "PATH", "LANG", "SSL_CERT_FILE", "OPENROUTER_API_KEY"}}
    files = {"/work/parts.json": json.dumps(describe(parts)),
             "/work/assessment.txt": assessment,
             "/work/schema.json": json.dumps(Refinement.model_json_schema()),
             "/work/cleanup-skill.md": instructions}
    if previous is not None:
        files["/work/previous-plan.json"] = json.dumps(previous["plan"])
        files["/work/previous-verdict.json"] = json.dumps(previous["verdict"])
        prompt += ("\nThis is a repair retry. Read /work/previous-plan.json and /work/previous-verdict.json. "
                   "Image order is reference, immutable original views, previous rejected candidate views. "
                   "The actual reviewer verdict explains the failure. Produce materially different supported "
                   "edits that address it, not the same plan with rewritten prose. All edits start from the "
                   "ORIGINAL parts, never cumulatively from the rejected candidate. If the defect requires "
                   "unsupported reconstruction, preserve source and explain it in unresolved.")
    encoded_images = [image_data_url(p) for p in images]
    feedback = ""
    for attempt in range(1, 4):
        process = subprocess.run([shutil.which("node") or "node", "--max-old-space-size=256",
                                  str(FX_RUNTIME / "run.mjs")],
                                 input=json.dumps({"prompt": prompt + feedback, "files": files,
                                                   "images": encoded_images}),
                                 text=True, capture_output=True, cwd=folder, env=env, timeout=260)
        if process.returncode:
            detail = (process.stderr + "\n" + process.stdout).replace(os.environ["OPENROUTER_API_KEY"], "[redacted]")
            (folder / "fx-error.txt").write_text(detail)
            raise RuntimeError(f"fx refinement failed (exit {process.returncode}): {detail[-1200:]}")
        response = json.loads(process.stdout)
        (folder / f"fx-evidence-{attempt}.json").write_text(json.dumps(response, indent=2))
        text = response.get("final_output") or response.get("output") or ""
        (folder / f"refinement-attempt-{attempt}.txt").write_text(text)
        try:
            plan = Refinement.model_validate(json_object(text))
            known = {key for key, _ in parts}
            edited = [edit.part_id for edit in plan.parts]
            if any(key not in known for key in edited) or len(edited) != len(set(edited)):
                raise ValueError("Use only listed component IDs, each at most once.")
        except ValueError as error:
            (folder / f"refinement-validation-{attempt}.txt").write_text(str(error))
            if attempt == 3:
                raise RuntimeError("fx did not produce a valid repair plan after three attempts; source retained.") from error
            feedback = (f"\nYour previous plan was invalid:\n{text}\nValidation errors:\n{error}\n"
                        "Return a corrected complete JSON plan. Preserve all bounds in the schema.")
            continue
        (folder / "refinement.json").write_text(plan.model_dump_json(indent=2))
        return plan
    raise AssertionError("Unreachable refinement retry state")


def review_candidate(brief, reference, before_views, after_views, plan):
    verdict = json_object(vision(
        "Compare before and after against the reference and brief. Image order: reference, before contact "
        "sheet with individual component previews, after assembled views. Judge only matched assembled views. "
        "Return JSON only with accept (boolean), quality_passed (boolean), unresolved (list of strings), "
        "blocking_issues (list of strings), repairable_with_supported_tools (boolean), and reason (string). "
        "Accept if the after asset has no visible regression in identity, silhouette, texture fidelity or "
        "materials. Independently set quality_passed only if the final asset is usable and faithful to the "
        "reference, with no blocking defects: broken anatomy, major identity loss, obvious texture corruption "
        "or seams. Minor stylization differences may remain unresolved without blocking. Do not excuse visible "
        "bad textures. A no-change or non-regressing result can still fail quality. Name remaining defects. "
        "Supported repair tools are bounded per-part RGB gains, roughness/metallic adjustment, duplicate or "
        "degenerate face cleanup, and visibility-aware front reference texture projection. Set "
        "repairable_with_supported_tools=true only when a different supported plan could address the actual "
        "failure. Missing or broken anatomy, substantial geometry reconstruction, unseen texture repainting "
        "and different reference pose are NOT repairable with these tools. Do not reward changes for their own sake. "
        f"Applied plan: {plan.model_dump_json()} Brief: {brief}", [reference, before_views, after_views]))
    if (type(verdict.get("accept")) is not bool or type(verdict.get("quality_passed")) is not bool
            or type(verdict.get("repairable_with_supported_tools")) is not bool
            or not isinstance(verdict.get("reason"), str)
            or not isinstance(verdict.get("unresolved"), list)
            or any(not isinstance(item, str) for item in verdict["unresolved"])
            or not isinstance(verdict.get("blocking_issues"), list)
            or any(not isinstance(item, str) for item in verdict["blocking_issues"])):
        raise ValueError("Astra returned an invalid post-refinement verdict.")
    verdict["model"] = MODEL
    verdict["unresolved"] = list(dict.fromkeys([*plan.unresolved, *verdict["unresolved"]]))
    verdict["quality_passed"] = verdict["accept"] and verdict["quality_passed"] and not verdict["blocking_issues"]
    verdict["source_retained"] = not verdict["quality_passed"]
    return verdict


def effective_edits(plan):
    data = plan.model_dump(exclude={"summary", "unresolved"})
    data["parts"] = sorted([{key: value for key, value in edit.items() if key != "name"}
                            for edit in data["parts"]], key=lambda edit: edit["part_id"])
    if not data["project_reference_texture"]:
        data.pop("reference_front_axis")
        data.pop("projection_strength")
    return data


def refine(parts, brief, reference, folder):
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    before = folder / "upscaled.glb"
    export(parts, before)
    before_views = render(before, folder / "before", components=True)
    assessment = vision("Inspect this generated asset against the brief. First image is its reference; second is "
                        "a labeled contact sheet of the assembled asset and components. Identify only visible "
                        "defects and name the affected part IDs. Preserve identity and texture channel meaning. "
                        f"Brief: {brief}", [reference, before_views])
    (folder / "visual-review.txt").write_text(assessment)
    previous = None
    previous_views = None
    previous_edits = None
    selected_parts = parts
    selected_views = before_views
    selected_attempt = None
    for number in range(1, 3):
        attempt = folder / f"repair-attempt-{number}"
        attempt.mkdir()
        images = [reference, before_views] + ([previous_views] if previous_views else [])
        plan = plan_with_fx(brief, parts, assessment, attempt, images=images, previous=previous)
        (attempt / "refinement.json").write_text(plan.model_dump_json(indent=2))
        edits = effective_edits(plan)
        if previous is not None and edits == previous_edits:
            verdict = dict(previous["verdict"], accept=False, quality_passed=False, source_retained=True,
                           reason="The retry repeated the rejected edits; source retained for review.",
                           stop_reason="identical_plan")
            (attempt / "verdict.json").write_text(json.dumps(verdict, indent=2))
            break
        candidate = apply_refinement(parts, plan)
        if plan.project_reference_texture:
            from .texture_repair import project_reference
            candidate, projection = project_reference(
                candidate, reference, attempt / "texture-projection",
                front_axis=plan.reference_front_axis, strength=plan.projection_strength)
            (attempt / "texture-projection.json").write_text(json.dumps(projection, indent=2))
        export(candidate, attempt / "candidate.glb")
        after_views = render(attempt / "candidate.glb", attempt / "after")
        verdict = review_candidate(brief, reference, before_views, after_views, plan)
        (attempt / "verdict.json").write_text(json.dumps(verdict, indent=2))
        if verdict["quality_passed"]:
            selected_parts, selected_views, selected_attempt = candidate, after_views, number
            break
        if not verdict["repairable_with_supported_tools"]:
            verdict["stop_reason"] = "unsupported_repairs"
            break
        previous = {"plan": plan.model_dump(), "verdict": verdict}
        previous_views, previous_edits = after_views, edits
    verdict = dict(verdict, attempts=number, selected_attempt=selected_attempt,
                   status="accepted" if selected_attempt else "needs_review",
                   candidate_accept=verdict["accept"], accept=selected_attempt is not None,
                   source_retained=selected_attempt is None)
    if not selected_attempt:
        verdict["quality_passed"] = False
        verdict.setdefault("stop_reason", "repair_limit")
    (folder / "refinement.json").write_text(plan.model_dump_json(indent=2))
    (folder / "verdict.json").write_text(json.dumps(verdict, indent=2))
    export(selected_parts, folder / "candidate.glb")
    if selected_views.is_file():
        shutil.copytree(selected_views.parent, folder / "after", dirs_exist_ok=True)
    return selected_parts, plan, verdict