Spaces:
Running on Zero
Running on Zero
Download studio/refine.py from mantrakp/component-studio-reference: direct link, hf CLI and curl.
- Browser
- Download file 14.9 kB
-
https://huggingface.co/spaces/mantrakp/component-studio-reference/resolve/main/studio/refine.py
- Command line
-
hf download hf://spaces/mantrakp/component-studio-reference/studio/refine.py
-
curl -L -o refine.py https://huggingface.co/spaces/mantrakp/component-studio-reference/resolve/main/studio/refine.py
14.9 kB
| """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 | |