inference/scripts/inference_psd_quantized.py

"""Quantized inference for See-through full pipeline (layerdiff body -> head -> marigold depth -> PSD).

Supports NF4 (default, 4-bit) and bf16 (baseline) modes. HF repos are auto-selected
based on quant_mode. Builds pipelines directly without using inference_utils global singletons.

Usage (from repo root):
    python inference/scripts/inference_psd_quantized.py --srcp image.png --save_to_psd
    python inference/scripts/inference_psd_quantized.py --quant_mode none --no_group_offload
"""

import os.path as osp
import argparse
import sys
import os
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))

default_n_threads = 8
os.environ['OPENBLAS_NUM_THREADS'] = f"{default_n_threads}"
os.environ['MKL_NUM_THREADS'] = f"{default_n_threads}"
os.environ['OMP_NUM_THREADS'] = f"{default_n_threads}"

import json
import time
import cv2
import numpy as np
import torch
from PIL import Image

from modules.layerdiffuse.diffusers_kdiffusion_sdxl import KDiffusionStableDiffusionXLPipeline
from modules.layerdiffuse.vae import TransparentVAE
from modules.layerdiffuse.layerdiff3d import UNetFrameConditionModel
from modules.marigold import MarigoldDepthPipeline
from utils.cv import center_square_pad_resize, smart_resize, img_alpha_blending
from utils.torch_utils import seed_everything
from utils.io_utils import json2dict, dict2json
from utils.inference_utils import further_extr
from utils.cv import validate_resolution


VALID_BODY_PARTS_V2 = [
    'hair', 'headwear', 'face', 'eyes', 'eyewear', 'ears', 'earwear', 'nose', 'mouth',
    'neck', 'neckwear', 'topwear', 'handwear', 'bottomwear', 'legwear', 'footwear',
    'tail', 'wings', 'objects'
]


def build_layerdiff_pipeline(args):
    """Build the LayerDiff3D pipeline with appropriate quantization."""
    quant_mode = args.quant_mode

    if quant_mode == 'none':
        # bf16 baseline: load from original repo
        repo = args.repo_id_layerdiff
        trans_vae = TransparentVAE.from_pretrained(repo, subfolder='trans_vae')
        unet = UNetFrameConditionModel.from_pretrained(repo, subfolder='unet')
        pipeline = KDiffusionStableDiffusionXLPipeline.from_pretrained(
            repo, trans_vae=trans_vae, unet=unet, scheduler=None)
        if args.cpu_offload:
            pipeline.vae.to(dtype=torch.bfloat16)
            pipeline.trans_vae.to(dtype=torch.bfloat16)
            pipeline.unet.to(dtype=torch.bfloat16)
            pipeline.text_encoder.to(dtype=torch.bfloat16)
            pipeline.text_encoder_2.to(dtype=torch.bfloat16)
            pipeline.enable_model_cpu_offload()
        else:
            pipeline.vae.to(dtype=torch.bfloat16, device='cuda')
            pipeline.trans_vae.to(dtype=torch.bfloat16, device='cuda')
            pipeline.unet.to(dtype=torch.bfloat16, device='cuda')
            pipeline.text_encoder.to(dtype=torch.bfloat16, device='cuda')
            pipeline.text_encoder_2.to(dtype=torch.bfloat16, device='cuda')
            if getattr(args, 'group_offload', False):
                pipeline.enable_group_offload('cuda', num_blocks_per_group=1)
        # Cache tag embeddings and unload text encoders to save VRAM
        pipeline.cache_tag_embeds()
    else:
        # NF4: load from pre-quantized repo (auto-selected by REPO_MAP)
        repo = args.repo_id_layerdiff
        unet = UNetFrameConditionModel.from_pretrained(repo, subfolder='unet')

        trans_vae = TransparentVAE.from_pretrained(repo, subfolder='trans_vae')  # always bf16
        pipeline = KDiffusionStableDiffusionXLPipeline.from_pretrained(
            repo, trans_vae=trans_vae, unet=unet, scheduler=None)

        if args.cpu_offload:
            # VAE + TransparentVAE to bf16; quantized components handled by bnb
            pipeline.vae.to(dtype=torch.bfloat16)
            pipeline.trans_vae.to(dtype=torch.bfloat16)
            pipeline.enable_model_cpu_offload()
        else:
            pipeline.vae.to(dtype=torch.bfloat16, device='cuda')
            pipeline.trans_vae.to(dtype=torch.bfloat16, device='cuda')
            # Don't manually .to(cuda) quantized components -- bnb handles device placement
            if getattr(args, 'group_offload', False):
                pipeline.enable_group_offload('cuda', num_blocks_per_group=1)
        # Cache tag embeddings and unload text encoders to save VRAM
        pipeline.cache_tag_embeds()

    return pipeline


def build_marigold_pipeline(args):
    """Build the Marigold depth pipeline with appropriate quantization."""
    quant_mode = args.quant_mode

    if quant_mode == 'none':
        repo = args.repo_id_depth
        unet = UNetFrameConditionModel.from_pretrained(repo, subfolder='unet')
        marigold_pipe = MarigoldDepthPipeline.from_pretrained(repo, unet=unet)
        if args.cpu_offload:
            marigold_pipe.to(dtype=torch.bfloat16)
            marigold_pipe.enable_model_cpu_offload()
        else:
            marigold_pipe.to(device='cuda', dtype=torch.bfloat16)
            if getattr(args, 'group_offload', False):
                marigold_pipe.enable_group_offload('cuda', num_blocks_per_group=1)
        marigold_pipe.cache_tag_embeds()
    else:
        # NF4: load from pre-quantized repo (auto-selected by REPO_MAP)
        repo = args.repo_id_depth
        unet = UNetFrameConditionModel.from_pretrained(repo, subfolder='unet', torch_dtype=torch.bfloat16)

        marigold_pipe = MarigoldDepthPipeline.from_pretrained(repo, unet=unet, torch_dtype=torch.bfloat16)
        marigold_pipe.vae.to(device='cuda')
        marigold_pipe.unet.to(device='cuda')
        # Text encoder may be quantized (from pre-quantized repo) — only move device, not dtype
        if not getattr(marigold_pipe.text_encoder, 'is_quantized', False) and \
           not getattr(marigold_pipe.text_encoder, 'quantization_method', None):
            marigold_pipe.text_encoder.to(device='cuda')
        if getattr(args, 'group_offload', False):
            marigold_pipe.enable_group_offload('cuda', num_blocks_per_group=1)
        marigold_pipe.cache_tag_embeds()

    return marigold_pipe


def run_layerdiff(pipeline, imgp, save_dir, seed, num_inference_steps, resolution):
    """Run LayerDiff3D body + head passes. Replicates inference_utils.py v3 logic exactly."""
    saved = osp.join(save_dir, osp.splitext(osp.basename(imgp))[0])
    os.makedirs(saved, exist_ok=True)
    input_img = np.array(Image.open(imgp).convert('RGBA'))
    fullpage, pad_size, pad_pos = center_square_pad_resize(input_img, resolution, return_pad_info=True)
    scale = pad_size[0] / resolution
    Image.fromarray(fullpage).save(osp.join(saved, 'src_img.png'))

    rng = torch.Generator(device=pipeline.unet.device).manual_seed(seed)

    # Body pass
    body_tag_list = ['front hair', 'back hair', 'head', 'neck', 'neckwear', 'topwear', 'handwear', 'bottomwear', 'legwear', 'footwear', 'tail', 'wings', 'objects']
    pipeline_output = pipeline(
        strength=1.0,
        num_inference_steps=num_inference_steps,
        batch_size=1,
        generator=rng,
        guidance_scale=1.0,
        prompt=body_tag_list,
        negative_prompt='',
        fullpage=fullpage,
        group_index=0
    )
    images = pipeline_output.images
    for rst, tag in zip(pipeline_output.images, body_tag_list):
        Image.fromarray(rst).save(osp.join(saved, f'{tag}.png'))
    head_img = images[2]

    # Head crop
    head_tag_list = ['headwear', 'face', 'irides', 'eyebrow', 'eyewhite', 'eyelash', 'eyewear', 'ears', 'earwear', 'nose', 'mouth']
    hx0, hy0, hw, hh = cv2.boundingRect(cv2.findNonZero((head_img[..., -1] > 15).astype(np.uint8)))

    hx = int(hx0 * scale) - pad_pos[0]
    hy = int(hy0 * scale) - pad_pos[1]
    hw = int(hw * scale)
    hh = int(hh * scale)

    def _crop_head(img, xywh):
        x, y, w, h = xywh
        ih, iw = img.shape[:2]
        x1 = x
        y1 = y
        x2 = x + w
        y2 = y + h
        if w < iw // 2:
            px = min(iw - x - w, x, w // 5)
            x1 = min(max(x - px, 0), iw)
            x2 = min(max(x + w + px, 0), iw)
        if h < ih // 2:
            py = min(ih - y - h, y, h // 5)
            y2 = min(max(y + h + py, 0), ih)
            y1 = min(max(y - py, 0), ih)
        return img[y1: y2, x1: x2], (x1, y1, x2, y2)

    input_head, (hx1, hy1, hx2, hy2) = _crop_head(input_img, [hx, hy, hw, hh])
    hx1 = int(hx1 / scale + pad_pos[0] / scale)
    hy1 = int(hy1 / scale + pad_pos[1] / scale)
    ih, iw = input_head.shape[:2]
    input_head, pad_size, pad_pos = center_square_pad_resize(input_head, resolution, return_pad_info=True)
    Image.fromarray(input_head).save(osp.join(saved, 'src_head.png'))

    # Head pass
    pipeline_output = pipeline(
        strength=1.0,
        num_inference_steps=num_inference_steps,
        batch_size=1,
        generator=rng,
        guidance_scale=1.0,
        prompt=head_tag_list,
        negative_prompt='',
        fullpage=input_head,
        group_index=1
    )
    canvas = np.zeros((resolution, resolution, 4), dtype=np.uint8)

    py1, py2, px1, px2 = (np.array([pad_pos[1], pad_pos[1] + ih, pad_pos[0], pad_pos[0] + iw]) / scale).astype(np.int64)

    scale_size = (int(pad_size[0] / scale), int(pad_size[1] / scale))

    for rst, tag in zip(pipeline_output.images, head_tag_list):
        rst = smart_resize(rst, scale_size)[py1: py2, px1: px2]
        full = canvas.copy()
        full[hy1: hy1 + rst.shape[0], hx1: hx1 + rst.shape[1]] = rst
        Image.fromarray(full).save(osp.join(saved, f'{tag}.png'))


def run_marigold(marigold_pipe, srcp, save_dir, seed, resolution_depth):
    """Run Marigold depth estimation. Matches inference_utils.apply_marigold logic.

    Uses resolution_depth to control Marigold inference resolution. If different from
    source image size, images are resized before depth prediction and depth maps are
    resized back after. All frames processed together (no chunking).
    """
    srcname = osp.basename(osp.splitext(srcp)[0])
    saved = osp.join(save_dir, srcname)

    # Read source image to get actual size (matches inference_utils approach)
    src_img_p = osp.join(saved, 'src_img.png')
    fullpage = np.array(Image.open(src_img_p).convert('RGBA'))
    src_h, src_w = fullpage.shape[:2]

    if isinstance(resolution_depth, int) and resolution_depth == -1:
        resolution_depth = [src_h, src_w]
    resolution_depth = validate_resolution(resolution_depth)
    src_rescaled = resolution_depth[0] != src_h or resolution_depth[1] != src_w

    img_list = []
    exist_list = []
    empty_array = np.zeros((src_h, src_w, 4), dtype=np.uint8)
    blended_alpha = np.zeros((src_h, src_w), dtype=np.float32)

    compose_list = {'eyes': ['eyewhite', 'irides', 'eyelash', 'eyebrow'], 'hair': ['back hair', 'front hair']}
    for tag in VALID_BODY_PARTS_V2:
        tagp = osp.join(saved, f'{tag}.png')
        if osp.exists(tagp):
            exist_list.append(True)
            tag_arr = np.array(Image.open(tagp))
            tag_arr[..., -1][tag_arr[..., -1] < 15] = 0
            img_list.append(tag_arr)
        else:
            img_list.append(empty_array)
            exist_list.append(False)

    compose_dict = {}
    for c, clist in compose_list.items():
        imlist = []
        taglist = []
        for tag in clist:
            p = osp.join(saved, tag + '.png')
            if osp.exists(p):
                tag_arr = np.array(Image.open(p))
                tag_arr[..., -1][tag_arr[..., -1] < 15] = 0
                imlist.append(tag_arr)
                taglist.append(tag)
        if len(imlist) > 0:
            img = img_alpha_blending(imlist, premultiplied=False)
            img_list[VALID_BODY_PARTS_V2.index(c)] = img
            compose_dict[c] = {'taglist': taglist, 'imlist': imlist}

    for img in img_list:
        blended_alpha += img[..., -1].astype(np.float32) / 255

    blended_alpha = np.clip(blended_alpha, 0, 1) * 255
    blended_alpha = blended_alpha.astype(np.uint8)
    fullpage[..., -1] = blended_alpha
    img_list.append(fullpage)

    # Resize to depth resolution if needed
    img_list_input = img_list
    if src_rescaled:
        img_list_input = [smart_resize(img, resolution_depth) for img in img_list]

    seed_everything(seed)
    pipe_out = marigold_pipe(color_map=None, img_list=img_list_input)
    depth_pred = pipe_out.depth_tensor
    depth_pred = depth_pred.to(device='cpu', dtype=torch.float32).numpy()

    # Resize depth back to source resolution if needed
    if src_rescaled:
        depth_pred = [smart_resize(d, (src_h, src_w)) for d in depth_pred]

    drawables = [{'img': img, 'depth': depth} for img, depth in zip(img_list, depth_pred)]
    drawables = drawables[:-1]
    blended = img_alpha_blending(drawables, premultiplied=False)

    infop = osp.join(saved, 'info.json')
    if osp.exists(infop):
        info = json2dict(infop)
    else:
        info = {'parts': {}}

    parts = info['parts']
    for ii, depth in enumerate(depth_pred[:-1]):
        depth = (np.clip(depth, 0, 1) * 255).astype(np.uint8)
        tag = VALID_BODY_PARTS_V2[ii]
        if tag in compose_dict:
            mask = blended_alpha > 256
            for t, im in zip(compose_dict[tag]['taglist'][::-1], compose_dict[tag]['imlist'][::-1]):
                mask_local = im[..., -1] > 15
                mask_invis = np.bitwise_and(mask, mask_local)
                depth_local = np.full((src_h, src_w), fill_value=255, dtype=np.uint8)
                depth_local[mask_local] = depth[mask_local]
                if np.any(mask_invis):
                    depth_local[mask_invis] = np.median(depth[np.bitwise_and(mask_local, np.bitwise_not(mask_invis))])
                mask = np.bitwise_or(mask, mask_local)

                parts_info = parts.get(t, {})
                Image.fromarray(depth_local).save(osp.join(saved, f'{t}_depth.png'))
                parts[t] = parts_info
            continue

        parts_info = parts.get(tag, {})
        Image.fromarray(depth).save(osp.join(saved, f'{tag}_depth.png'))
        parts[tag] = parts_info

    dict2json(info, infop)
    Image.fromarray(blended).save(osp.join(saved, 'reconstruction.png'))


if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description="Quantized inference: LayerDiff body+head -> Marigold depth -> PSD"
    )
    parser.add_argument('--srcp', type=str, default='assets/test_image.png', help='input image')
    parser.add_argument('--save_dir', type=str, default='workspace/layerdiff_output')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--resolution', type=int, default=1280)
    parser.add_argument('--save_to_psd', action='store_true')
    parser.add_argument('--tblr_split', action='store_true',
                        help='try split parts (handwear, eyes, etc) into left-right components')
    parser.add_argument('--quant_mode', type=str, default='nf4', choices=['nf4', 'none'],
                        help='quantization mode: nf4 (default, 4-bit) or none (bf16 baseline)')
    parser.add_argument('--repo_id_layerdiff', type=str, default=None,
                        help='Override LayerDiff3D HF repo (auto-selected based on quant_mode)')
    parser.add_argument('--repo_id_depth', type=str, default=None,
                        help='Override Marigold3D HF repo (auto-selected based on quant_mode)')
    parser.add_argument('--cpu_offload', action='store_true', default=False,
                        help='enable model CPU offload (default: on)')
    parser.add_argument('--no_cpu_offload', action='store_false', dest='cpu_offload',
                        help='disable model CPU offload')
    parser.add_argument('--num_inference_steps', type=int, default=30)
    parser.add_argument('--resolution_depth', type=int, default=768,
                        help='Marigold depth inference resolution (default 768; -1 to match layerdiff resolution)')
    parser.add_argument('--group_offload', action='store_true', default=True,
                        help='Enable group offload to reduce peak VRAM (default: on)')
    parser.add_argument('--no_group_offload', action='store_false', dest='group_offload',
                        help='Disable group offload for faster inference on high-VRAM GPUs')
    args = parser.parse_args()

    # Auto-select HF repos based on quant_mode
    REPO_MAP = {
        'nf4': {
            'layerdiff': '24yearsold/seethroughv0.0.2_layerdiff3d_nf4',
            'depth': '24yearsold/seethroughv0.0.1_marigold_nf4',
        },
        'none': {
            'layerdiff': 'layerdifforg/seethroughv0.0.2_layerdiff3d',
            'depth': '24yearsold/seethroughv0.0.1_marigold',
        },
    }
    defaults = REPO_MAP[args.quant_mode]
    if args.repo_id_layerdiff is None:
        args.repo_id_layerdiff = defaults['layerdiff']
    if args.repo_id_depth is None:
        args.repo_id_depth = defaults['depth']

    srcp = args.srcp
    seed = args.seed
    resolution = args.resolution
    num_inference_steps = args.num_inference_steps
    save_dir = args.save_dir
    srcname = osp.basename(osp.splitext(srcp)[0])
    saved = osp.join(save_dir, srcname)

    print(f"Quantized inference: quant_mode={args.quant_mode}, cpu_offload={args.cpu_offload}")
    print(f"  Source image: {srcp}")
    print(f"  Save dir: {save_dir}")
    print(f"  Resolution: {resolution}, Steps: {num_inference_steps}, Seed: {seed}")

    torch.cuda.reset_peak_memory_stats()
    total_t0 = time.time()

    # --- LayerDiff ---
    print('\nBuilding LayerDiff3D pipeline...')
    seed_everything(seed)
    pipeline = build_layerdiff_pipeline(args)

    print('Running LayerDiff3D (body + head)...')
    layerdiff_t0 = time.time()
    run_layerdiff(pipeline, srcp, save_dir, seed, num_inference_steps, resolution)
    layerdiff_time = time.time() - layerdiff_t0
    print(f'  LayerDiff3D done in {layerdiff_time:.1f}s')

    # Free layerdiff pipeline before loading marigold
    del pipeline
    torch.cuda.empty_cache()

    # --- Marigold ---
    print('\nBuilding Marigold depth pipeline...')
    marigold_pipe = build_marigold_pipeline(args)

    print('Running Marigold depth...')
    marigold_t0 = time.time()
    run_marigold(marigold_pipe, srcp, save_dir, seed, resolution_depth=args.resolution_depth)
    marigold_time = time.time() - marigold_t0
    print(f'  Marigold done in {marigold_time:.1f}s')

    # Free marigold pipeline before PSD assembly
    del marigold_pipe
    torch.cuda.empty_cache()

    # --- PSD assembly ---
    print('\nRunning PSD assembly...')
    psd_t0 = time.time()
    further_extr(saved, rotate=False, save_to_psd=args.save_to_psd, tblr_split=args.tblr_split)
    psd_time = time.time() - psd_t0
    print(f'  PSD assembly done in {psd_time:.1f}s')

    total_time = time.time() - total_t0

    # --- Stats ---
    stats = {
        'quant_mode': args.quant_mode,
        'peak_vram_gb': torch.cuda.max_memory_allocated() / 1024**3,
        'layerdiff_time_s': layerdiff_time,
        'marigold_time_s': marigold_time,
        'psd_time_s': psd_time,
        'total_time_s': total_time,
    }
    print(f'\n{"="*60}')
    print(json.dumps(stats, indent=2))
    print(f'{"="*60}')
    with open(osp.join(saved, 'stats.json'), 'w') as f:
        json.dump(stats, f, indent=2)
    print(f'Stats saved to {osp.join(saved, "stats.json")}')