"""Preload the released pipeline once, then run requests without model CLIs.""" from contextlib import contextmanager import importlib.util import json import os from pathlib import Path import sys import time from hf_space.model_pool import ModelPool ROOT = Path(__file__).resolve().parents[1] def _configure_preload_threads(torch, *, cpu_only): # spaces 0.51.3 already sets the process-wide inter-op count at import. # PyTorch permits that setter only once, even before the pool is consumed. # Intra-op is adjustable and stays serial for the later ZeroGPU fork. torch.set_num_threads(1 if cpu_only else 4) if not cpu_only and torch.get_num_interop_threads() != 1: torch.set_num_interop_threads(1) class ForegroundResident: """The standalone cutout task needs only RMBG, without native 3D modules.""" def __init__(self, model_path): self.config = {'rmbg_model': str(model_path)} self.pool = ModelPool(directory='/tmp/supratex-foreground-state-v1') def prepare(self, notify): import torch from transformers import AutoModelForImageSegmentation _configure_preload_threads(torch, cpu_only=True) notify('Loading background-removal model', 0, 1) self.pool.add('background', AutoModelForImageSegmentation.from_pretrained( self.config['rmbg_model'], trust_remote_code=True, local_files_only=True).eval()) self.pool.sealed = True notify('Background-removal model ready', 1, 1) class ResidentPipeline: def __init__(self, config): self.config_path = str(config) self.config = config self.pool = ModelPool() self.geometry = None self.qwen = None self.material_config = None self.material_step = None self.material_processor = None def prepare(self, notify, *, cpu_only=False): import torch from scripts.run_example import read_config from supratex.inference.runtime import configure_inference_runtime # Startup constructors run on host tensors; no model forward is called. config = read_config(Path(self.config)) self.config = config os.environ.update(TORCH_HOME=config['torch_home'], TOKENIZERS_PARALLELISM='false', ATTN_BACKEND='sdpa', SPARSE_ATTN_BACKEND='sdpa', SUPRATEX_RUNTIME_PROFILE='torch28', SUPRATEX_LOW_VRAM='1') # The isolated dedicated-GPU worker can globally disable Hub access. # ZeroGPU preloads in the HTTP process, which still serves OAuth and # Bucket requests; local model paths and local_files_only are sufficient. if not cpu_only: os.environ.update(HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1') _configure_preload_threads(torch, cpu_only=cpu_only) if not cpu_only and not torch.cuda.is_available(): raise RuntimeError('The GPU is unavailable; startup models cannot become ready') for path in (ROOT, ROOT/'deployment/runtime/moge', ROOT/'deployment/runtime/utils3d', Path(config['geometry_package']), Path(config['qwen_runtime']), Path(config['qwen_deps']) if config.get('qwen_deps') else None): if path is not None and str(path) not in sys.path: sys.path.insert(0, str(path)) configure_inference_runtime(attention='sdpa', low_vram=True, device='cpu' if cpu_only else 'cuda', material_backend='qwen-image') total = 18 done = 0 def load(name, label, factory): nonlocal done notify(label, done, total) model = self.pool.add(name, factory()) done += 1 notify(label, done, total) return model from transformers import AutoModelForImageSegmentation load('background', 'Loading background-removal model', lambda: AutoModelForImageSegmentation.from_pretrained(config['rmbg_model'], trust_remote_code=True, local_files_only=True).eval()) from moge.model.v2 import MoGeModel moge = Path(config['moge_model']) load('camera', 'Loading camera model', lambda: MoGeModel.from_pretrained(str(moge/'model.pt' if moge.is_dir() else moge)).eval()) from supratex.frontend.geometry_runner import build_native_pipeline notify('Loading geometry models', done, total) def geometry_loaded(name, model): label = ('Loading geometry image encoder' if '.conditioner.' in name else 'Loading geometry model: '+name.removeprefix('geometry.').replace('_', ' ')) return load(name, label, lambda: model) self.geometry = build_native_pipeline(Path(config['geometry_model']), model_loaded=geometry_loaded) # Keep native Qwen's local imports explicit and separate from root infer.py. native_dir = Path(config['qwen_runtime'])/'examples/qwen_image_21/pbr' sys.path.insert(0, str(native_dir)) spec = importlib.util.spec_from_file_location('_supratex_resident_qwen', native_dir/'infer.py') module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) self.qwen = module from supratex.frontend.qwen_image import validate _, base, adapter, _ = validate(Path(config['qwen_runtime']), config['material_base'], config['material_model'], deps=Path(config['qwen_deps']) if config.get('qwen_deps') else None, local_files_only=True) cfg, weights, step = module.load_checkpoint(adapter) self.material_config, self.material_step = cfg, step if cfg['deterministic_algorithms']: os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') os.environ['DIFFSYNTH_ATTENTION_IMPLEMENTATION'] = cfg['attention_backend'] from transformers import AutoProcessor self.material_processor = AutoProcessor.from_pretrained(str(base/'processor'), local_files_only=True) self.material_processor.image_processor.size = {'shortest_edge': cfg['resolution']**2, 'longest_edge': cfg['resolution']**2} from diffsynth.models.qwen_image_21_vae import QwenImage21VAE from diffsynth.models.qwen_image_21_text_encoder import QwenImage21TextEncoder from diffsynth.utils.state_dict_converters.qwen_image_21_text_encoder import QwenImage21TextEncoderStateDictConverter load('material.vae', 'Loading material VAE', lambda: module.explicit_load( QwenImage21VAE, module.checkpoint_component(base, 'vae'), 'cpu', torch.bfloat16)) load('material.text', 'Loading material image-and-text encoder', lambda: module.explicit_load( QwenImage21TextEncoder, module.checkpoint_component(base, 'text_encoder'), 'cpu', torch.bfloat16, QwenImage21TextEncoderStateDictConverter)) def material_model(): model = module.build_model(cfg | {'model_root': str(base)}, 'cpu') module.strict_load_adapter(model, weights) return model.eval() load('material.transformer', 'Loading material transformer and adapter', material_model) del weights from supratex.inference.pipeline import validate_completion_artifacts, completion_decoder_stems, _completion_vae_root, _load_vae from supratex import models from supratex.inference.checkpoint import load_completion_state_dict from supratex.pipelines.observation_guided_completion import build_completion_model c, dataset, checkpoint, _ = validate_completion_artifacts(config['model_config'], config['checkpoint'], config['trellis_root']) ckpts = _completion_vae_root(config['trellis_root']) vaes = {'shape_encoder': ckpts/'shape_enc_next_dc_f16c32_fp16', 'texture_encoder': ckpts/'tex_enc_next_dc_f16c32_fp16'} stems = completion_decoder_stems(c, dataset, config['model_config'], config['trellis_root']) vaes.update(shape_decoder=stems['pretrained_shape_slat_dec'], material_decoder=stems['pretrained_pbr_slat_dec']) assert total == done + len(vaes) + 1 for name, stem in vaes.items(): load('completion.'+name, 'Loading completion '+name.replace('_', ' '), lambda stem=stem: _load_vae(models, stem).eval()) load('completion.transformer', 'Loading completion transformer', lambda: build_completion_model(c, load_completion_state_dict(checkpoint, torch_module=torch), assign_weights=True).eval()) # Establish a working CUDA context before enabling Generate. Constructors # above load host weights; they do not need neural warm-up forwards. if not cpu_only: probe = torch.zeros(1, device='cuda') torch.cuda.synchronize() del probe torch.cuda.empty_cache() self.pool.sealed = True notify('Models ready', total, total) print('[resident-models] '+json.dumps(self.pool.summary()), flush=True) @contextmanager def use_geometry(self): try: # Native geometry already moves its individual networks by stage. yield self.geometry finally: self.pool.restore_all() def run_geometry(self, **kwargs): from supratex.frontend.geometry_runner import run with self.use_geometry() as pipeline: return run(**kwargs, pipeline=pipeline) def run_material(self, input_path, output, *, device): from hf_space.timing import span from diffsynth.pipelines.qwen_image_21 import QwenImage21Pipeline, QwenImage21Unit_EditImageEmbedder, QwenImage21Unit_PromptEmbedder import torch m, pool = self.qwen, self.pool cfg = self.material_config original, square, transform = m.prepare_image(input_path, cfg['resolution']) output.mkdir(parents=True, exist_ok=False) previous_deterministic = torch.are_deterministic_algorithms_enabled() previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() if cfg['deterministic_algorithms']: os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8') started = time.monotonic() device = torch.device('cuda:0' if str(device) == 'cuda' else device) torch.cuda.set_device(device) pipe = QwenImage21Pipeline(device=device, torch_dtype=torch.bfloat16) pipe.processor = self.material_processor vae, encoder = pool.get('material.vae'), pool.get('material.text') model = pool.get('material.transformer') try: torch.use_deterministic_algorithms(cfg['deterministic_algorithms']) with span('qwen.conditions'): pipe.vae = vae with m.stage_on_device(vae, device, enabled=True), torch.inference_mode(): edited = QwenImage21Unit_EditImageEmbedder().process(pipe, edit_image=[square], height=cfg['resolution'], width=cfg['resolution'], tiled=False, tile_size=256, tile_stride=192) if edited['edit_image'][0].size != square.size: raise ValueError('native edit path changed dimensions') rgb = edited['edit_latents'][0] pipe.text_encoder = encoder blocks = m.text_encoder_offload_blocks(encoder) with m.sequential_cpu_offload(encoder, blocks, device, enabled=True), torch.inference_mode(): values = QwenImage21Unit_PromptEmbedder().process(pipe, prompt=m.PROMPT, edit_image=[square]) conditions = m.canonical_conditions(dict(values, rgb=rgb), cfg['resolution']) pipe.text_encoder = None del values, edited, blocks, rgb with span('qwen.predict_decode'): latent, decoded = m.predict_and_decode(model, vae, conditions, cfg, low_vram=True) with span('qwen.output_write'): maps, _ = m.output_maps(decoded, square, original, transform, False) for name, image in zip(('albedo', 'mra'), maps): image.save(output/(name+'.png')) (output/'REPORT.json').write_text(json.dumps({ 'status': 'complete', 'checkpoint_step': self.material_step, 'config': cfg, 'model_loading': 'startup_pool', 'transform': transform, 'low_vram': True, 'memory_policy': 'cpu_stages_and_sequential_transformer_blocks', 'elapsed_seconds': time.monotonic()-started}, indent=2)) finally: # Restore immutable host mappings even when an inference hook fails. try: pool.restore_all() finally: torch.use_deterministic_algorithms(previous_deterministic, warn_only=previous_warn_only) def completion_model(self, name): return self.pool.get('completion.'+name)