#!/usr/bin/env python3 """Generate a textured PBR mesh from one RGB/RGBA image.""" from __future__ import annotations import argparse import json import math import os from pathlib import Path import sys def _pipeline_overrides(args: argparse.Namespace) -> dict: """Validate explicit model selections before resolving a release.""" overrides = {} for component, attribute in (("completion", "pretrained"), ("material", "material_model"), ("qwen-base", "material_base"), ("geometry", "geometry_model"), ("vae", "trellis_root"), ("moge2", "moge_model"), ("rmbg2", "rmbg_model")): value = getattr(args, attribute) if value is not None: if component not in ("completion", "material") and not Path(value).expanduser().exists(): raise ValueError(f"--pipeline requires a local --{attribute.replace('_', '-')} override") overrides[component] = value if args.revision is not None: if args.pretrained is None: raise ValueError("--revision requires --pretrained with --pipeline; use --pipeline-revision for the release") overrides["completion_revision"] = args.revision if args.material_revision is not None: if args.material_model is None: raise ValueError("--material-revision requires --material-model with --pipeline") overrides["material_revision"] = args.material_revision return overrides def _prepare_pipeline(args: argparse.Namespace, overrides: dict, descriptor: dict, *, has_alpha: bool) -> dict: """Prepare the release already selected during argument validation.""" from scripts.prepare_models import prepare_pipeline from scripts.run_example import read_config source = descriptor["repo_id"] if descriptor["kind"] == "hub" else args.pipeline revision = descriptor.get("revision") if descriptor["kind"] == "hub" else args.pipeline_revision prepared = prepare_pipeline(source, output_dir=args.models_dir, revision=revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only, with_rmbg=not has_alpha, overrides=overrides) return read_config(prepared) def _validate_inference_options(args: argparse.Namespace) -> None: """Check explicit values first, then call again after release defaults.""" from scripts.prepare_models import _pipeline if args.preview_only and (args.dual_view or args.refine): raise ValueError("o-voxel preview is single-view and precedes mesh refinement/export") if args.back_view_dir is not None and not args.dual_view: raise ValueError("--back-view-dir requires --dual-view") _pipeline.validate_export_options(args) if not math.isfinite(args.geometry_initial_mesh_scale) or args.geometry_initial_mesh_scale <= 0: raise ValueError("geometry-initial-mesh-scale must be finite and positive") if not 1 <= args.geometry_target_padding <= 30 or args.geometry_max_iterations < 1: raise ValueError("geometry-target-padding must be in [1, 30]; geometry-max-iterations must be positive") if not 0 <= args.resolution_level <= 9: raise ValueError("resolution-level must be in [0, 9]") def _validate_input_output(args: argparse.Namespace, *, require_rmbg: bool = True) -> tuple[Path, Path, Path, bool]: """Check local inputs and destinations without creating or downloading files.""" image, output = args.input.expanduser().resolve(), args.output.expanduser().absolute() expected_suffix = ".npz" if args.preview_only else ".glb" if output.suffix.lower() != expected_suffix: raise ValueError(f"--output must have a {expected_suffix} suffix") if output.exists() or output.is_symlink(): raise FileExistsError(f"refusing to overwrite existing output: {output}") output = output.resolve() work = (args.work_dir or output.parent / f"{output.stem}_work").expanduser().absolute() if work.is_symlink() and not work.exists(): raise FileExistsError(f"--work-dir must be a fresh empty directory: {work}") work = work.resolve() if work.exists() and (not work.is_dir() or any(work.iterdir())): raise FileExistsError(f"--work-dir must be a fresh empty directory: {work}") for destination in (output, work): for parent in destination.parents: if parent.exists(): if not parent.is_dir(): raise NotADirectoryError(f"destination parent must be a directory: {parent}") break if not image.is_file(): raise FileNotFoundError(f"input image is missing: {image}") from PIL import Image with Image.open(image) as source: low, high = source.getchannel("A").getextrema() if source.mode == "RGBA" else (255, 255) has_alpha = low < 255 and high > 0 if args.use_input_alpha and not has_alpha: raise ValueError("--use-input-alpha requires a non-opaque RGBA input with foreground pixels") if require_rmbg and not has_alpha and not args.rmbg_model: raise ValueError("RGB/opaque input requires --rmbg-model") source.load() if args.dual_view and args.back_view_dir is None: import shutil if shutil.which("codex") is None: raise RuntimeError("automatic back generation requires Codex with gpt-6-astra and built-in image_gen") return image, output, work, has_alpha def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False) parser.add_argument("--pipeline", help="SUPRA-Tex release directory, pipeline.json, or HF repo ID; prepares and loads all stages") parser.add_argument("--pipeline-revision", help="HF revision of the complete SUPRA-Tex release") parser.add_argument("--models-dir", type=Path, help="prepared model directory for --pipeline (default: ckpts/ready)") parser.add_argument("--input", type=Path, help="source RGB/RGBA image (default: assets/spiny_whelk_shell.png, or " "assets/cinema_signage.png with --dual-view)") parser.add_argument("--output", required=True, type=Path, help="new destination .glb file") parser.add_argument("--preview-only", action="store_true", help="save/render final o-voxel; export mesh separately") parser.add_argument("--work-dir", type=Path, help="fresh intermediate directory (default: _work beside the GLB)") parser.add_argument("--moge-model", help="local MoGe-2 checkpoint/directory or HF repo ID") parser.add_argument("--rmbg-model", help="local RMBG-2 directory or HF repo ID; required for RGB input") parser.add_argument("--use-input-alpha", action="store_true", help="require and reuse foreground RGBA alpha") parser.add_argument("--no-filter", "--no-image-depth", dest="no_filter", action="store_true", help="disable the default single-view image/geometry depth pre-filter; dual-view does not use it") parser.add_argument("--resolution-level", type=int, default=9, help="MoGe-2 resolution level") parser.add_argument("--geometry-package", type=Path, help="native or paper geometry source checkout; see docs/geometry.md") parser.add_argument("--geometry-model", help="local geometry weights directory or HF repo ID") parser.add_argument("--geometry-revision", help="geometry weights HF revision") parser.add_argument("--geometry-initial-mesh-scale", type=float, default=1.0) parser.add_argument("--geometry-scale-search", action="store_true", help="enable bounded mesh-scale calibration (default: Pixal3D fixed scale)") parser.add_argument("--parallel-frontend", action="store_true", help="run geometry and Qwen material concurrently on the first two visible GPUs") parser.add_argument("--geometry-target-padding", type=int, default=3, help="scale-search target padding") parser.add_argument("--geometry-max-iterations", type=int, default=2, help="scale-search update limit") parser.add_argument("--firered-package", type=Path, default=Path(__file__).resolve().parent / "supratex/frontend/material_model") parser.add_argument("--material-backend", choices=("qwen-image", "firered"), default="qwen-image") parser.add_argument("--material-base", help="selected backend base model: local directory or HF repo ID") parser.add_argument("--firered-base", help="legacy FireRed base option; requires --material-backend firered") parser.add_argument("--qwen-runtime", type=Path, help="native Qwen source snapshot (default: bundled runtime)") parser.add_argument("--qwen-deps", type=Path, help="optional Qwen dependency overlay; otherwise use the active Python environment") parser.add_argument("--material-model", help="local backend checkpoint bundle or HF repo ID") parser.add_argument("--material-revision", help="material checkpoint HF revision") parser.add_argument("--material-weights", type=Path) parser.add_argument("--material-spec", type=Path) parser.add_argument("--material-checkpoint-manifest", type=Path) parser.add_argument("--pretrained", help="local Supra-Tex completion bundle or HF repo ID") parser.add_argument("--revision", help="completion bundle HF revision") parser.add_argument("--model-config", type=Path, help="completion ABI JSON; overrides bundle config") parser.add_argument("--checkpoint", type=Path, help="completion weights; overrides bundle checkpoint") parser.add_argument("--trellis-root", type=Path, help="TRELLIS.2 VAE directory (all four components for HF releases)") parser.add_argument("--cache-dir", type=Path, help="shared HF download cache for all adapters") parser.add_argument("--local-files-only", action="store_true", help="resolve models locally without Hub downloads") parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), help="inference attention backend; auto selects by GPU and installed extensions") parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, help="CPU offload for Qwen and 3D stages; auto enabled on GPUs with up to 24 GiB") parser.add_argument("--seed", type=int, help="geometry and completion seed (material runner uses its fixed seed 42)") parser.add_argument("--texture-size", type=int) parser.add_argument("--decimation-target", type=int) parser.add_argument("--texture-format", choices=("png", "webp")) parser.add_argument("--gutter", action="store_true", help="enable chart-isolated one-texel base-color gutter padding (default: disabled); " "with --refine, the observation-guided gutter is used") parser.add_argument("--refine", action="store_true", help="refine the visible texture after export: write the observed reference-view albedo into the atlas " "and repair its UV gutters from the observation (sharper fine text and detail; needs the 4096 PNG " "atlas; adds about 2 minutes on an RTX 3090)") parser.add_argument("--remove-floaters", action=argparse.BooleanOptionalAction, default=True, help="before completion, drop disconnected geometry pieces that project mostly outside the " "image foreground (default: on); extra geometry attached to the object is kept") parser.add_argument("--dual-view", action="store_true", help="generate a back view (needs the external Codex CLI) and texture with both views; includes --refine for both views") parser.add_argument("--back-view-dir", type=Path, help="reuse a prepared, geometry-bound back package instead of generating a new back") args = parser.parse_args(argv) from scripts.run_example import MODEL_KEYS, _pipeline_config, default_input, inference_defaults if args.input is None: args.input = default_input(args.dual_view) if not args.pipeline and (args.pipeline_revision is not None or args.models_dir is not None): parser.error("--pipeline-revision and --models-dir require --pipeline") material_paths = (args.material_weights, args.material_spec, args.material_checkpoint_manifest) if args.pipeline and args.material_backend != "qwen-image": parser.error("--pipeline requires --material-backend qwen-image") if args.material_backend == "qwen-image" and (args.firered_base or any(material_paths)): parser.error("FireRed options require --material-backend firered") try: # Reject explicit argument/input errors before fetching any metadata. _validate_inference_options(args) if args.pipeline: overrides = _pipeline_overrides(args) _, _, _, has_alpha = _validate_input_output(args, require_rmbg=False) from scripts.prepare_models import _pipeline manifest, descriptor = _pipeline.resolve_pipeline( args.pipeline, revision=args.pipeline_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only) defaults = manifest["defaults"] else: defaults = inference_defaults() for key, value in defaults.items(): if getattr(args, key) is None: setattr(args, key, value) # Release defaults can introduce a dual-view conflict; check before downloads. _validate_inference_options(args) if args.pipeline: prepared = _prepare_pipeline(args, overrides, descriptor, has_alpha=has_alpha) except (OSError, ValueError, RuntimeError) as error: parser.error(str(error)) if args.pipeline: explicit_pretrained = args.pretrained is not None direct_overrides = {"geometry_package", "qwen_runtime", "qwen_deps", "model_config", "checkpoint"} for key in MODEL_KEYS: if prepared.get(key) is not None and (key not in direct_overrides or getattr(args, key) is None): if explicit_pretrained and key in ("model_config", "checkpoint"): continue action = parser._option_string_actions["--" + key.replace("_", "-")] value = prepared[key] setattr(args, key, action.type(value) if action.type else value) args.local_files_only = True args._pipeline_config = prepared if not all((args.moge_model, args.geometry_package, args.geometry_model)): parser.error("provide --pipeline or all of --moge-model, --geometry-package, and --geometry-model") if args.material_backend == "qwen-image": args.qwen_runtime = args.qwen_runtime or _pipeline_config().runtime_root(Path(__file__).resolve().parent) / "qwen" args.material_base = args.material_base or str(Path(__file__).resolve().parent / "ckpts/qwen-base") args.material_model = args.material_model or str(Path(__file__).resolve().parent / "ckpts/qwen-material") else: if args.qwen_runtime is not None or args.qwen_deps is not None: parser.error("--qwen-runtime and --qwen-deps require --material-backend qwen-image") if args.material_base and args.firered_base: parser.error("provide only one of --material-base and --firered-base") if not (args.material_base or args.firered_base): parser.error("FireRed requires --material-base or --firered-base") if args.material_model and any(material_paths): parser.error("--material-model cannot be combined with explicit material artifact paths") if not args.material_model and not all(material_paths): parser.error("provide --material-model or all three --material-weights, --material-spec, --material-checkpoint-manifest") if not args.pretrained and not all((args.model_config, args.checkpoint, args.trellis_root)): parser.error("provide --pretrained or all of --model-config, --checkpoint, and --trellis-root") return args def _completion_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]: """Resolve local/Hub artifacts and check missing completion files before GPU work.""" from supratex.weights import COMPLETION_ALLOW_PATTERNS, resolve_completion_artifacts, resolve_model_dir root = None if args.pretrained: root = resolve_model_dir(args.pretrained, revision=args.revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only, allow_patterns=COMPLETION_ALLOW_PATTERNS) return resolve_completion_artifacts(root, model_config=args.model_config, checkpoint=args.checkpoint, trellis_root=args.trellis_root) def run(args: argparse.Namespace, *, resident=None, foreground_prepared: bool = False) -> Path: """Run the existing image frontend, then complete and export its generated mesh.""" if resident is not None and (args.material_backend != "qwen-image" or args.dual_view or args.refine or args.parallel_frontend): raise ValueError("Resident inference supports sequential single-view Qwen generation without refinement") if type(foreground_prepared) is not bool or (foreground_prepared and resident is None): raise ValueError('Prepared foreground is reserved for the trusted resident pipeline') if getattr(args, "_pipeline_config", None) is not None: from scripts.run_example import runtime_environment environment = runtime_environment(args._pipeline_config, args.attention) for key in ("ATTN_BACKEND", "SPARSE_ATTN_BACKEND"): if key not in environment: os.environ.pop(key, None) os.environ.update(environment) for source in reversed(environment["PYTHONPATH"].split(os.pathsep)): if source not in sys.path: sys.path.insert(0, source) image, output, work, _ = _validate_input_output(args) from supratex.inference.runtime import configure_inference_runtime runtime = configure_inference_runtime(attention=getattr(args, "attention", "auto"), low_vram=getattr(args, "low_vram", None), material_backend=args.material_backend) if resident is not None and not runtime["low_vram"]: raise ValueError("Resident inference requires the released low-VRAM stage policy") # The adapters validate source compatibility and full checkpoint contents. from supratex.frontend.geometry import validate_source validate_source(args.geometry_package.expanduser()) if args.material_backend == "qwen-image": from supratex.frontend.qwen_image import validate validate(args.qwen_runtime, args.material_base, args.material_model, deps=args.qwen_deps, revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only) elif not (args.firered_package.expanduser() / "infer.py").is_file(): raise FileNotFoundError(f"material entrypoint is missing: {args.firered_package / 'infer.py'}") material_paths = {} for name in ("material_weights", "material_spec", "material_checkpoint_manifest"): path = getattr(args, name) if path is not None: path = path.expanduser().resolve() if not path.is_file(): raise FileNotFoundError(f"missing {name}: {path}") material_paths[name] = path config, checkpoint, encoders = _completion_paths(args) from supratex.frontend.pipeline import run as run_frontend from supratex.inference import run_completion work.mkdir(parents=True, exist_ok=True) handoff = work / "input" filter_requested = not getattr(args, 'no_filter', False) filter_enabled = filter_requested and not args.dual_view filter_disabled_reason = ('user_disabled' if not filter_requested else 'dual_view_unsupported' if args.dual_view else None) receipt = {"status": "running", "stage": "frontend", "input": str(image), "output": str(output), "work_dir": str(work), "seed": args.seed, "gutter": args.gutter, "material_backend": args.material_backend, "material_seed": 42, "runtime": runtime, "model_config": str(config), "checkpoint": str(checkpoint), "dual_view": args.dual_view, "refine": args.refine, "remove_floaters": args.remove_floaters, "preview_only": args.preview_only, "foreground_prepared": foreground_prepared, "image_depth_filter": {"requested": filter_requested, "enabled": filter_enabled, "disabled_reason": filter_disabled_reason}} def record() -> None: temporary = work / ".run.json.tmp" temporary.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") temporary.replace(work / "run.json") record() try: print(f"[1/2] Foreground, MoGe-2, geometry, material maps -> {handoff}", flush=True) result = run_frontend(image, handoff, moge_model=args.moge_model, rmbg_model=args.rmbg_model, use_input_alpha=args.use_input_alpha, resolution_level=args.resolution_level, export_image_depth=filter_enabled, foreground_prepared=foreground_prepared, geometry_package=args.geometry_package.expanduser().resolve(), geometry_model=args.geometry_model, geometry_revision=args.geometry_revision, geometry_seed=args.seed, geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, geometry_scale_search=args.geometry_scale_search, parallel_frontend=args.parallel_frontend, geometry_target_padding=args.geometry_target_padding, geometry_max_iterations=args.geometry_max_iterations, firered_package=args.firered_package.expanduser().resolve(), firered_base=args.firered_base, material_backend=args.material_backend, material_base=args.material_base, qwen_runtime=args.qwen_runtime if args.material_backend == "qwen-image" else None, qwen_deps=args.qwen_deps if args.material_backend == "qwen-image" else None, material_model=args.material_model, **material_paths, revision=args.material_revision, cache_dir=args.cache_dir, local_files_only=args.local_files_only, **({"resident": resident} if resident is not None else {}), **({"record_material_provenance": True} if args.dual_view else {})) optional = {"refine": True, "work_dir": work / "refine"} if args.refine else {} if args.dual_view: receipt.update(stage="back_view", frontend=result) record() if args.back_view_dir is None: from supratex.backview.workflow import build_back_view back_dir = build_back_view(handoff, work / "back_view", cache_dir=args.cache_dir, local_files_only=args.local_files_only) else: back_dir = args.back_view_dir.expanduser().resolve() optional = {"back_view_dir": back_dir, "work_dir": work / "dual_view"} receipt["back_view_dir"] = str(back_dir) receipt.update(stage="completion", frontend=result) record() print(f"[2/2] Surface observations, PBR completion, {'o-voxel preview' if args.preview_only else 'GLB export'} -> {output}", flush=True) run_completion(input_dir=handoff, output=output, model_config=config, checkpoint=checkpoint, trellis_root=encoders, seed=args.seed, texture_size=args.texture_size, decimation_target=args.decimation_target, texture_format=args.texture_format, gutter=args.gutter, remove_floaters=args.remove_floaters, preview_only=args.preview_only, **({"resident": resident} if resident is not None else {}), **optional) if not output.is_file() or output.stat().st_size == 0: raise RuntimeError(f"completion did not emit a nonempty result: {output}") except Exception as error: receipt.update(status="failed", error=f"{type(error).__name__}: {error}") record() raise finally: if resident is not None: resident.pool.restore_all() receipt.update(status="complete", stage="complete") record() print(f"Textured mesh: {output}\nIntermediate files and run receipt: {work}", flush=True) return output def main() -> None: run(parse_args()) if __name__ == "__main__": main()