""" 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)