import argparse import json import os from pathlib import Path from typing import Dict, Optional, List from omegaconf import OmegaConf from PIL import Image from dataclasses import dataclass from collections import defaultdict import torch import torch.utils.checkpoint from torchvision.utils import make_grid 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 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 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_tensor_image(tensor, fp): ndarr = convert_to_numpy(tensor) save_image_numpy(ndarr, fp) return ndarr def save_image_numpy(ndarr, fp): im = Image.fromarray(ndarr) im.save(fp) @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_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 # new two-stage settings run_mode: str = "full" # full | generate | reconstruct multiview_tmp_dir: str = "" prefer_edited_views: bool = True save_multiview_metadata: bool = True def ensure_rgba(img: Image.Image) -> Image.Image: return img.convert("RGBA") if img.mode != "RGBA" else img def get_scene_name(batch, sample_index: int) -> str: return Path(batch['filename'][sample_index]).stem def get_scene_dir(base_dir: str, scene: str) -> Path: return Path(base_dir) / scene def save_multiview_scene(base_dir: str, scene: str, colors: List[Image.Image], normals: List[Image.Image], meta: Optional[dict] = None): scene_dir = get_scene_dir(base_dir, scene) raw_dir = scene_dir / "raw" edit_dir = scene_dir / "edit" raw_dir.mkdir(parents=True, exist_ok=True) edit_dir.mkdir(parents=True, exist_ok=True) for idx, img in enumerate(colors): img = ensure_rgba(img) img.save(raw_dir / f"color_{idx:02d}.png") img.save(edit_dir / f"color_{idx:02d}.png") for idx, img in enumerate(normals): img = ensure_rgba(img) img.save(raw_dir / f"normal_{idx:02d}.png") img.save(edit_dir / f"normal_{idx:02d}.png") if meta is not None: with open(scene_dir / "meta.json", "w", encoding="utf-8") as f: json.dump(meta, f, indent=2) def load_multiview_scene(base_dir: str, scene: str, prefer_edit=True): scene_dir = get_scene_dir(base_dir, scene) candidate_dirs = [scene_dir / ("edit" if prefer_edit else "raw"), scene_dir / ("raw" if prefer_edit else "edit")] data_dir = None for cdir in candidate_dirs: if cdir.exists(): data_dir = cdir break if data_dir is None: raise FileNotFoundError(f"No multiview directory found for scene '{scene}' under {scene_dir}") color_paths = sorted(data_dir.glob("color_*.png")) normal_paths = sorted(data_dir.glob("normal_*.png")) if not color_paths or not normal_paths: raise FileNotFoundError(f"No color/normal images found in {data_dir}") colors = [ensure_rgba(Image.open(p)) for p in color_paths] normals = [ensure_rgba(Image.open(p)) for p in normal_paths] return colors, normals def prepare_scene_views(batch, imgs_in, normals_pred, images_pred, out, cfg: TestConfig, save_dir, images_cond, case_id): guidance_scale = cfg.validation_guidance_scales num_views = imgs_in.shape[0] // (out.shape[0] // 2 // cfg.num_views) if False else None # unused safeguard bsz = out.shape[0] // 2 num_views = cfg.num_views scene_results = [] if cfg.save_mode == 'concat': 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 = get_scene_name(batch, i) img_in_ = images_cond[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_tensor_image(vis_, out_filename) return scene_results if cfg.save_mode != 'rgb': raise ValueError(f"Unsupported save_mode for two-stage workflow: {cfg.save_mode}") for i in range(bsz // num_views): scene = get_scene_name(batch, i) normals, colors = [], [] for j in range(num_views): idx = i * num_views + j normal = normals_pred[idx] if j == 0: color = imgs_in[i * num_views].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) normals[0][:, :256, 256:512] = normals[-1] color_pils = [ensure_rgba(remove(convert_to_pil(tensor), session=session)) for tensor in colors[:6]] normal_pils = [ensure_rgba(remove(convert_to_pil(tensor), session=session)) for tensor in normals[:6]] meta = None if cfg.save_multiview_metadata: meta = { "scene": scene, "case_id": case_id, "num_colors": len(color_pils), "num_normals": len(normal_pils), "seed": cfg.seed, "run_mode": cfg.run_mode, "crop_size": cfg.validation_dataset.crop_size, "with_smpl": cfg.with_smpl, } scene_results.append((scene, color_pils, normal_pils, meta)) return scene_results def run_inference(dataloader, econdata, pipeline, carving, cfg: TestConfig, save_dir): if pipeline is not None: pipeline.set_progress_bar_config(disable=True) if cfg.seed is None: generator = None else: device = pipeline.unet.device if pipeline is not None else "cuda" generator = torch.Generator(device=device).manual_seed(cfg.seed) for case_id, batch in tqdm(enumerate(dataloader)): if cfg.run_mode == "reconstruct": batch_size = len(batch['filename']) for i in range(batch_size): scene = get_scene_name(batch, i) colors, normals = load_multiview_scene( cfg.multiview_tmp_dir, scene, prefer_edit=cfg.prefer_edited_views, ) pose = econdata.__getitem__(case_id + i) carving.optimize_case(scene, pose, colors, normals) torch.cuda.empty_cache() continue images_cond = 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") 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 = batch['normal_prompt_embeddings'] clr_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"): 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:] scene_results = prepare_scene_views( batch=batch, imgs_in=imgs_in, normals_pred=normals_pred, images_pred=images_pred, out=out, cfg=cfg, save_dir=save_dir, images_cond=images_cond, case_id=case_id, ) if cfg.save_mode == 'concat': continue for i, (scene, colors, normals, meta) in enumerate(scene_results): if cfg.run_mode == "generate": save_multiview_scene(cfg.multiview_tmp_dir, scene, colors, normals, meta=meta) print(f"[PSHuman] Saved multiview scene '{scene}' to {get_scene_dir(cfg.multiview_tmp_dir, scene)}") continue pose = econdata.__getitem__(case_id + i) 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(): pipeline.to('cuda') return pipeline def main(cfg: TestConfig): if cfg.seed is not None: set_seed(cfg.seed) pipeline = None if cfg.run_mode == "reconstruct" else 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 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) if cfg.run_mode in {"generate", "reconstruct"} and not cfg.multiview_tmp_dir: raise ValueError("multiview_tmp_dir must be provided for run_mode='generate' or 'reconstruct'.") 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 cfg = load_config(args.config, cli_args=extras) schema = OmegaConf.structured(TestConfig) cfg = OmegaConf.merge(schema, cfg) main(cfg)