inference/scripts/inference_psd_hypersd.py

"""
HyperSD-accelerated LayerDiff inference.

Applies HyperSD SDXL 4-step LoRA to the LayerDiff UNet, reducing inference
from 30 steps to 4. Marigold depth estimation is unchanged.

Usage:
    python inference/scripts/inference_psd_hypersd.py \
        --srcp assets/test_image.png --save_to_psd

    # With group offload for lower VRAM
    python inference/scripts/inference_psd_hypersd.py \
        --srcp assets/test_image.png --save_to_psd --group_offload
"""

import os.path as osp
import argparse
import sys
import os
import gc

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 numpy as np
import torch
from tqdm import tqdm
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
from diffusers import DDIMScheduler

from utils.io_utils import find_all_imgs
from utils import inference_utils
from utils.inference_utils import apply_marigold, further_extr
from utils.torch_utils import seed_everything

from modules.layerdiffuse.diffusers_kdiffusion_sdxl import KDiffusionStableDiffusionXLPipeline, UNetFrameConditionModel
from modules.layerdiffuse.vae import TransparentVAE


HYPERSD_REPO = "ByteDance/Hyper-SD"
LORA_FILES = {
    1: "Hyper-SDXL-1step-lora.safetensors",
    2: "Hyper-SDXL-2steps-lora.safetensors",
    4: "Hyper-SDXL-4steps-lora.safetensors",
    8: "Hyper-SDXL-8steps-lora.safetensors",
}


def apply_lora_to_unet(unet, lora_state_dict, alpha_scale=1.0):
    """
    Manually merge LoRA weights into UNet parameters in-place.

    For each LoRA layer: W_new = W + (alpha / rank) * lora_up @ lora_down

    Works without peft dependency by directly modifying model weights.
    """
    # Build lookup: underscore-joined module path -> actual module path
    module_lookup = {}
    for name, _ in unet.named_modules():
        key = name.replace('.', '_')
        module_lookup[key] = name

    # Group LoRA keys by module
    lora_groups = {}
    for key in lora_state_dict:
        # Keys are like: lora_unet_down_blocks_0_resnets_0_conv1.lora_down.weight
        # or: unet.down_blocks.0.resnets.0.conv1.lora_down.weight (PEFT format)
        if '.lora_down.weight' in key:
            module_key = key.replace('.lora_down.weight', '')
        elif '.lora_up.weight' in key:
            module_key = key.replace('.lora_up.weight', '')
        elif '.alpha' in key:
            module_key = key.replace('.alpha', '')
        else:
            continue

        if module_key not in lora_groups:
            lora_groups[module_key] = {}

        if '.lora_down.weight' in key:
            lora_groups[module_key]['down'] = lora_state_dict[key]
        elif '.lora_up.weight' in key:
            lora_groups[module_key]['up'] = lora_state_dict[key]
        elif '.alpha' in key:
            lora_groups[module_key]['alpha'] = lora_state_dict[key]

    applied = 0
    skipped = 0

    for module_key, lora_data in lora_groups.items():
        if 'down' not in lora_data or 'up' not in lora_data:
            continue

        # Convert LoRA key to module path
        # Strip common prefixes
        clean_key = module_key
        for prefix in ['lora_unet_', 'unet_', 'unet.']:
            if clean_key.startswith(prefix):
                clean_key = clean_key[len(prefix):]
                break

        # Try to find the module
        # First try direct lookup (underscore format)
        module_path = module_lookup.get(clean_key)

        if module_path is None:
            # Try dot format (PEFT-style keys)
            clean_dot = clean_key.replace('_', '.')
            # Check if this exact path exists
            try:
                module = unet
                for part in clean_dot.split('.'):
                    module = getattr(module, part)
                module_path = clean_dot
            except (AttributeError, IndexError):
                module_path = None

        if module_path is None:
            skipped += 1
            continue

        # Get the module and its weight
        module = unet
        try:
            for part in module_path.split('.'):
                if part.isdigit():
                    module = module[int(part)]
                else:
                    module = getattr(module, part)
        except (AttributeError, IndexError, TypeError):
            skipped += 1
            continue

        if not hasattr(module, 'weight'):
            skipped += 1
            continue

        down = lora_data['down'].to(device=module.weight.device, dtype=module.weight.dtype)
        up = lora_data['up'].to(device=module.weight.device, dtype=module.weight.dtype)
        alpha = lora_data.get('alpha', torch.tensor(down.shape[0], dtype=torch.float32))
        alpha = alpha.item() if isinstance(alpha, torch.Tensor) else alpha
        rank = down.shape[0]
        scale = (alpha / rank) * alpha_scale

        # Compute delta: reshape to 2D, multiply, reshape back
        orig_shape = module.weight.shape
        if down.ndim == 4:
            # Conv LoRA: down=[rank, in, kh, kw], up=[out, rank, 1, 1]
            down_2d = down.reshape(rank, -1)
            up_2d = up.reshape(up.shape[0], rank)
            delta = (up_2d @ down_2d).reshape(orig_shape)
        else:
            # Linear LoRA: down=[rank, in], up=[out, rank]
            delta = up @ down

        if delta.shape != orig_shape:
            skipped += 1
            continue

        module.weight.data += scale * delta
        applied += 1

    print(f"LoRA merge: {applied} layers applied, {skipped} layers skipped")
    return applied


def build_hypersd_pipeline(args):
    """Build LayerDiff pipeline with HyperSD LoRA applied."""
    pretrained = args.repo_id_layerdiff
    lora_steps = args.lora_steps

    print(f"Loading LayerDiff pipeline from {pretrained}...")
    trans_vae = TransparentVAE.from_pretrained(pretrained, subfolder='trans_vae')
    unet = UNetFrameConditionModel.from_pretrained(pretrained, subfolder='unet')
    pipeline = KDiffusionStableDiffusionXLPipeline.from_pretrained(
        pretrained, trans_vae=trans_vae, unet=unet, scheduler=None
    )

    # Download and apply HyperSD LoRA
    lora_filename = LORA_FILES[lora_steps]
    print(f"Downloading HyperSD LoRA: {lora_filename}...")
    lora_path = hf_hub_download(HYPERSD_REPO, lora_filename)

    print("Loading LoRA weights...")
    lora_sd = load_file(lora_path)

    # Try diffusers load_lora_weights first, fall back to manual merge
    try:
        pipeline.load_lora_weights(lora_sd)
        pipeline.fuse_lora()
        print("LoRA applied via diffusers load_lora_weights + fuse_lora")
    except Exception as e:
        print(f"load_lora_weights failed ({e}), falling back to manual LoRA merge...")
        apply_lora_to_unet(pipeline.unet, lora_sd)

    # Swap scheduler to DDIM with trailing timestep spacing (required by HyperSD)
    pipeline.scheduler = DDIMScheduler.from_config(
        pipeline.scheduler.config,
        timestep_spacing="trailing"
    )
    print(f"Scheduler swapped to DDIMScheduler (trailing, {lora_steps} steps)")

    # Move to GPU
    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 args.group_offload:
        pipeline.enable_group_offload('cuda', num_blocks_per_group=1)

    return pipeline


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="HyperSD-accelerated LayerDiff inference")
    parser.add_argument('--save_dir', type=str, default='workspace/layerdiff_output')
    parser.add_argument('--srcp', type=str, default='assets/test_image.png', help='input image or directory')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--repo_id_layerdiff', default='layerdifforg/seethroughv0.0.2_layerdiff3d')
    parser.add_argument('--repo_id_depth', default='24yearsold/seethroughv0.0.1_marigold')
    parser.add_argument('--resolution', type=int, default=1280, help="inference resolution of layerdiff")
    parser.add_argument('--resolution_depth', type=int, default=720, help="inference resolution of depth model")
    parser.add_argument('--lora_steps', type=int, default=4, choices=[1, 2, 4, 8],
                        help="HyperSD LoRA variant (determines step count)")
    parser.add_argument('--num_steps', type=int, default=None,
                        help="override inference steps (default: match lora_steps)")
    parser.add_argument('--save_to_psd', action='store_true')
    parser.add_argument('--tblr_split', action='store_true')
    parser.add_argument('--disable_progressbar', action='store_true')
    parser.add_argument('--group_offload', action='store_true')
    args = parser.parse_args()

    if args.num_steps is None:
        args.num_steps = args.lora_steps

    srcp = args.srcp
    if osp.isdir(srcp):
        imglist = find_all_imgs(srcp, abs_path=True)
    else:
        imglist = [srcp]

    # Build pipeline with HyperSD LoRA
    pipeline = build_hypersd_pipeline(args)

    # Inject into inference_utils so apply_layerdiff reuses it
    inference_utils.layerdiff_pipeline = pipeline

    for srcp in tqdm(imglist):
        seed_everything(args.seed)

        print(f'running layerdiff (HyperSD {args.lora_steps}-step)...')
        inference_utils.apply_layerdiff(
            srcp, args.repo_id_layerdiff,
            save_dir=args.save_dir, seed=args.seed,
            resolution=args.resolution,
            disable_progressbar=args.disable_progressbar,
            num_inference_steps=args.num_steps,
            group_offload=args.group_offload
        )

        print('running marigold...')
        apply_marigold(
            srcp, args.repo_id_depth,
            save_dir=args.save_dir, seed=args.seed,
            disable_progressbar=args.disable_progressbar,
            resolution=args.resolution_depth,
            group_offload=args.group_offload
        )

        srcname = osp.basename(osp.splitext(srcp)[0])
        saved = osp.join(args.save_dir, srcname)
        further_extr(saved, rotate=False, save_to_psd=args.save_to_psd, tblr_split=args.tblr_split)