| |
| |
| """Experimental, staged Idol inference for Apple MPS. |
| |
| Upstream revision: 4702a7f00850e34b6e83a2170cc123c3e9570744. |
| No weights are bundled. Full-weight inference has NOT been verified by this project. |
| """ |
| import gc |
| import json |
| import os |
| import shutil |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| from .engine import device_for |
|
|
|
|
| def patch_attention(layers, query_chunk=128): |
| """Bound attention working memory by splitting only the query dimension. |
| |
| Each query still sees ALL keys/values. Does not split or change softmax's key axis. |
| Disable CUDA-specific optimized sampling and avoid complex RoPE tensors on MPS. |
| """ |
| if query_chunk < 1: |
| raise ValueError("query_chunk must be positive") |
| from torch.nn import functional as F |
|
|
| def attention(*args): |
| q, k, v = args[0].unbind(dim=2) if len(args) == 1 else args |
| q, k, v = (t.permute(0, 2, 1, 3) for t in (q, k, v)) |
| chunks = [F.scaled_dot_product_attention(q[:, :, start:start+query_chunk], k, v) |
| for start in range(0, q.shape[2], query_chunk)] |
| return torch.cat(chunks, dim=2).permute(0, 2, 1, 3) |
|
|
| layers.scaled_dot_product_attention = attention |
| layers.RePo3DRotaryEmbedding.real_valued = True |
| return attention |
|
|
|
|
| def release(device): |
| gc.collect() |
| if device.type == "mps": |
| torch.mps.synchronize() |
| torch.mps.empty_cache() |
|
|
|
|
| @torch.no_grad() |
| def run_idol(args): |
| if args.steps < 1 or args.gaussians < 32768 or args.gaussians > 1048576 or args.gaussians % 32: |
| raise ValueError("Steps must be positive; Idol requires 32768..1048576 Gaussians, divisible by 32") |
| device = device_for(args.device) |
| cache = Path(args.cache).resolve() |
| cache.mkdir(parents=True, exist_ok=True) |
| os.environ["HF_HOME"] = str(cache) |
| os.environ["HF_HUB_OFFLINE"] = "0" if args.allow_download else "1" |
| os.environ["ATTN_BACKEND"] = "sdpa" |
| if args.allow_download and shutil.disk_usage(cache).free < 16 * 1024**3: |
| raise ValueError("Less than 16 GiB disk space available. This conservative download guard leaves room for teacher weights, encoders and temporary files. Free space, or use an existing complete cache without --allow-download.") |
| try: |
| from huggingface_hub import hf_hub_download, snapshot_download |
| from safetensors.torch import load_file |
| from tensordict import TensorDict |
| from idol.conditioning import ImageEncoder, require_encoder_access |
| from idol.models import BiRefNet, GaussianFixedlenDecoder, LatentSeqMMFlowModel, OctreeProbabilityFixedlenDecoder, layers |
| from idol.preprocess import CROP_MARGIN, preprocess_image, has_real_alpha |
| from idol.sampler import sample_flow |
| from idol.camera import decode_camera |
| except ImportError as exc: |
| raise RuntimeError("Install the optional teacher dependencies: python -m pip install '.[teacher]'") from exc |
| out = Path(args.out) |
| if out.exists(): |
| raise ValueError("Teacher output directory already exists; choose a new --out") |
| |
| config_path = hf_hub_download(args.weights, "config.json", local_files_only=not args.allow_download) |
| config = json.loads(Path(config_path).read_text()) |
| encoders = config["image_encoders"] |
| require_encoder_access([encoders["dinov3"], encoders["flux2_vae"]]) |
| weights = Path(snapshot_download(args.weights, allow_patterns=["config.json", "*.safetensors"], |
| local_files_only=not args.allow_download)) |
| patch_attention(layers, args.query_chunk) |
| torch.manual_seed(args.seed) |
|
|
| def load(model, name): |
| state = load_file(str(weights / f"{name}.safetensors"), device="cpu") |
| model.load_state_dict(state, strict=True) |
| del state |
| return model.eval().requires_grad_(False).to(device) |
|
|
| print("Experimental Idol MPS adapter: full-weight compatibility and quality remain unverified.", flush=True) |
| with Image.open(args.image) as opened: |
| image = opened.copy() |
| background = None |
| if not has_real_alpha(image): |
| print("Stage 1: background removal", flush=True) |
| background = load(BiRefNet(), "birefnet") |
| prepared = preprocess_image(image, background, config["image_size"], CROP_MARGIN[args.mode]) |
| del background |
| release(device) |
| pixels = torch.from_numpy(np.array(prepared)).permute(2, 0, 1).float()[None].div(255).to(device) |
| print("Stage 2: image conditioning", flush=True) |
| encoder = ImageEncoder(encoders["dinov3"], encoders["flux2_vae"], device="cpu") |
| |
| encoder.dinov3.to(device) |
| feature1 = encoder._dinov3_tokens(pixels) |
| del encoder.dinov3 |
| release(device) |
| encoder.flux2_vae.to(device) |
| feature2 = encoder._flux2_tokens(pixels) |
| del encoder, pixels |
| release(device) |
| cond = TensorDict({"feature1": feature1, "feature2": feature2}, batch_size=1, device=device) |
| print("Stage 3: latent sampling", flush=True) |
| model = load(LatentSeqMMFlowModel(**config["denoiser"]["args"]), "denoiser") |
| noise = TensorDict({"latent": torch.randn(1, model.q_token_length, model.in_channels, device=device), |
| "camera": torch.randn(1, 1, model.cam_channels, device=device)}, batch_size=1, device=device) |
| result = sample_flow(model, noise, cond=cond, neg_cond=cond.apply(torch.zeros_like), steps=args.steps, |
| rescale_t=3., cfg_strength=3., verbose=True, batch_cfg=False) |
| latent, camera = result["latent"], decode_camera(result["camera"][0, 0]) |
| del model, noise, cond, feature1, feature2, result |
| release(device) |
| print("Stage 4: octree decoding", flush=True) |
| model = load(OctreeProbabilityFixedlenDecoder(**config["decoder"]["args"]), "decoder") |
| per_anchor = config["decoder_gs"]["args"]["representation_config"]["num_gaussians"] |
| anchors = OctreeProbabilityFixedlenDecoder.sample(model, latent, num_points=args.gaussians // per_anchor, |
| level=config["max_voxel_level"]) |
| del model |
| release(device) |
| print("Stage 5: Gaussian decoding", flush=True) |
| model = load(GaussianFixedlenDecoder(**config["decoder_gs"]["args"]), "decoder_gs") |
| gaussian = model.to_representation(anchors, model(anchors, latent))[0] |
| out.mkdir(parents=True) |
| gaussian.save_ply(str(out / "teacher.ply")) |
| prepared.save(out / "input.png") |
| (out / "metadata.json").write_text(json.dumps({"teacher": args.weights, "steps": args.steps, |
| "gaussians": args.gaussians, "seed": args.seed, "device": str(device), "camera": camera, |
| "sdk_revision": "4702a7f00850e34b6e83a2170cc123c3e9570744", "weights_snapshot": weights.name, |
| "coordinate_frame": "Idol exported PLY (y-up)", "adapter": "experimental staged MPS"}, indent=2)) |
| print(f"Teacher output saved to {out}", flush=True) |
|
|
