inference.py

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)