import argparse import sys import os # spaces/ZeroGPU sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True by default, which # hits an NVML assertion (CUDACachingAllocator.cpp:1131) on MIG-sliced GPUs. Disable it # before torch initializes its CUDA allocator. os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:False" from typing import Dict, Optional, Tuple, List from omegaconf import OmegaConf from PIL import Image from dataclasses import dataclass from collections import defaultdict import torch if torch.cuda.is_available(): _props = torch.cuda.get_device_properties(0) print(f"GPU_DIAG: {_props.name} total_memory={_props.total_memory/1e9:.1f}GB", flush=True) import torch.utils.checkpoint from torchvision.utils import make_grid, save_image from accelerate.utils import set_seed from tqdm.auto import tqdm import torch.nn.functional as F from einops import rearrange from rembg import remove, new_session import pdb from mvdiffusion.pipelines.pipeline_mvdiffusion_unclip import StableUnCLIPImg2ImgPipeline from econdataset import SMPLDataset from reconstruct import ReMesh providers = [ ('CUDAExecutionProvider', { 'device_id': 0, 'arena_extend_strategy': 'kSameAsRequested', 'gpu_mem_limit': 8 * 1024 * 1024 * 1024, 'cudnn_conv_algo_search': 'HEURISTIC', }) ] session = new_session(providers=providers) weight_dtype = torch.float16 def tensor_to_numpy(tensor): return tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() @dataclass class TestConfig: pretrained_model_name_or_path: str revision: Optional[str] validation_dataset: Dict save_dir: str seed: Optional[int] validation_batch_size: int dataloader_num_workers: int # save_single_views: bool save_mode: str local_rank: int pipe_kwargs: Dict pipe_validation_kwargs: Dict unet_from_pretrained_kwargs: Dict validation_guidance_scales: float validation_grid_nrow: int num_views: int enable_xformers_memory_efficient_attention: bool with_smpl: Optional[bool] recon_opt: Dict def convert_to_numpy(tensor): return tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() def convert_to_pil(tensor): return Image.fromarray(convert_to_numpy(tensor)) def save_image(tensor, fp): ndarr = convert_to_numpy(tensor) # pdb.set_trace() save_image_numpy(ndarr, fp) return ndarr def save_image_numpy(ndarr, fp): im = Image.fromarray(ndarr) im.save(fp) def run_inference(dataloader, econdata, pipeline, carving, cfg: TestConfig, save_dir): pipeline.set_progress_bar_config(disable=True) if cfg.seed is None: generator = None else: generator = torch.Generator(device=pipeline.unet.device).manual_seed(cfg.seed) images_cond, pred_cat = [], defaultdict(list) for case_id, batch in tqdm(enumerate(dataloader)): images_cond.append(batch['imgs_in'][:, 0]) imgs_in = torch.cat([batch['imgs_in']]*2, dim=0) num_views = imgs_in.shape[1] imgs_in = rearrange(imgs_in, "B Nv C H W -> (B Nv) C H W")# (B*Nv, 3, H, W) if cfg.with_smpl: smpl_in = torch.cat([batch['smpl_imgs_in']]*2, dim=0) smpl_in = rearrange(smpl_in, "B Nv C H W -> (B Nv) C H W") else: smpl_in = None normal_prompt_embeddings, clr_prompt_embeddings = batch['normal_prompt_embeddings'], batch['color_prompt_embeddings'] prompt_embeddings = torch.cat([normal_prompt_embeddings, clr_prompt_embeddings], dim=0) prompt_embeddings = rearrange(prompt_embeddings, "B Nv N C -> (B Nv) N C") with torch.autocast("cuda"): # B*Nv images guidance_scale = cfg.validation_guidance_scales unet_out = pipeline( imgs_in, None, prompt_embeds=prompt_embeddings, dino_feature=None, smpl_in=smpl_in, generator=generator, guidance_scale=guidance_scale, output_type='pt', num_images_per_prompt=1, **cfg.pipe_validation_kwargs ) out = unet_out.images bsz = out.shape[0] // 2 normals_pred = out[:bsz] images_pred = out[bsz:] if cfg.save_mode == 'concat': ## save concatenated color and normal--------------------- pred_cat[f"cfg{guidance_scale:.1f}"].append(torch.cat([normals_pred, images_pred], dim=-1)) # b, 3, h, w cur_dir = os.path.join(save_dir, f"cropsize-{cfg.validation_dataset.crop_size}-cfg{guidance_scale:.1f}-seed{cfg.seed}-smpl-{cfg.with_smpl}") os.makedirs(cur_dir, exist_ok=True) for i in range(bsz//num_views): scene = batch['filename'][i].split('.')[0] img_in_ = images_cond[-1][i].to(out.device) vis_ = [img_in_] for j in range(num_views): idx = i*num_views + j normal = normals_pred[idx] color = images_pred[idx] vis_.append(color) vis_.append(normal) out_filename = f"{cur_dir}/{scene}.png" vis_ = torch.stack(vis_, dim=0) vis_ = make_grid(vis_, nrow=len(vis_), padding=0, value_range=(0, 1)) save_image(vis_, out_filename) elif cfg.save_mode == 'rgb': for i in range(bsz//num_views): scene = batch['filename'][i].split('.')[0] img_in_ = images_cond[-1][i].to(out.device) normals, colors = [], [] for j in range(num_views): idx = i*num_views + j normal = normals_pred[idx] if j == 0: color = imgs_in[0].to(out.device) else: color = images_pred[idx] if j in [3, 4]: normal = torch.flip(normal, dims=[2]) color = torch.flip(color, dims=[2]) colors.append(color) if j == 6: normal = F.interpolate(normal.unsqueeze(0), size=(256, 256), mode='bilinear', align_corners=False).squeeze(0) normals.append(normal) ## save color and normal--------------------- # normal_filename = f"normals_{view}_masked.png" # rgb_filename = f"color_{view}_masked.png" # save_image(normal, os.path.join(scene_dir, normal_filename)) # save_image(color, os.path.join(scene_dir, rgb_filename)) normals[0][:, :256, 256:512] = normals[-1] colors = [remove(convert_to_pil(tensor), session=session) for tensor in colors[:6]] normals = [remove(convert_to_pil(tensor), session=session) for tensor in normals[:6]] pose = econdata.__getitem__(case_id) carving.optimize_case(scene, pose, colors, normals) torch.cuda.empty_cache() def load_pshuman_pipeline(cfg): pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(cfg.pretrained_model_name_or_path, torch_dtype=weight_dtype) pipeline.unet.enable_xformers_memory_efficient_attention() if torch.cuda.is_available(): # A blanket .to('cuda') OOMs on this ZeroGPU MIG slice's effective VRAM. # MMGP dynamically offloads unused submodules to CPU RAM and quantizes the # main model, instead of keeping the whole pipeline resident in VRAM at once. from mmgp import offload, profile_type # image_normalizer uses a custom .scale() method rather than .forward(), so # mmgp's just-in-time weight-loading hook never fires for it, leaving its # mean/std buffers stranded on CPU. It's two tiny tensors, so just pin it on # GPU directly instead of routing it through mmgp's offload management. mmgp_components = {k: v for k, v in pipeline.components.items() if k != "image_normalizer"} offload.profile(mmgp_components, profile_type.VerylowRAM_LowVRAM, verboseLevel=1) if getattr(pipeline, "image_normalizer", None) is not None: pipeline.image_normalizer.to("cuda") # offload.profile() leaves torch'''s global default device set to '''cuda'''. # mmgp'''s own hooks move their tracked models''' weights explicitly regardless of # this setting, but plenty of *other* code later in this script (econdataset, # pymaf, smplx) constructs plain tensors with no explicit device, assuming the # normal CPU default. Reset it so that code behaves as originally written. torch.set_default_device('''cpu''') return pipeline def main( cfg: TestConfig ): # If passed along, set the training seed now. if cfg.seed is not None: set_seed(cfg.seed) pipeline = load_pshuman_pipeline(cfg) if cfg.with_smpl: from mvdiffusion.data.testdata_with_smpl import SingleImageDataset else: from mvdiffusion.data.single_image_dataset import SingleImageDataset # Get the dataset validation_dataset = SingleImageDataset( **cfg.validation_dataset ) validation_dataloader = torch.utils.data.DataLoader( validation_dataset, batch_size=cfg.validation_batch_size, shuffle=False, num_workers=cfg.dataloader_num_workers ) dataset_param = {'image_dir': validation_dataset.root_dir, 'seg_dir': None, 'colab': False, 'has_det': True, 'hps_type': 'pixie'} econdata = SMPLDataset(dataset_param, device='cuda') carving = ReMesh(cfg.recon_opt, econ_dataset=econdata) run_inference(validation_dataloader, econdata, pipeline, carving, cfg, cfg.save_dir) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, required=True) args, extras = parser.parse_known_args() from utils.misc import load_config # parse YAML config to OmegaConf cfg = load_config(args.config, cli_args=extras) schema = OmegaConf.structured(TestConfig) cfg = OmegaConf.merge(schema, cfg) main(cfg)