inference.py

# -------------------------------------------
#
# File Name: inference.py
# Author: WANG Yiyang
# Created: October.7, 2025
# Description: Command-line interface for DiffCamera for inference.
# -------------------------------------------
import argparse
import copy
import cv2
import os
from peft import LoraConfig, set_peft_model_state_dict
from peft.utils import get_peft_model_state_dict
from diffusers import (
    AutoencoderKL,
    FlowMatchEulerDiscreteScheduler,
)
from diffusers.utils.torch_utils import is_compiled_module
from src.flux.pipeline.pipeline_flux_new import FluxPipeline, flux_pipeline_call
from src.flux.model.transformer_flux_new import FluxTransformer2DModel
from accelerate import Accelerator
import torch
import torch.nn as nn
from utils import parse_args
import safetensors
from diffusers.utils import (
    check_min_version,
    convert_unet_state_dict_to_peft,
    is_wandb_available,
)
from PIL import Image
from torchvision import transforms
import numpy as np
from tqdm import tqdm

args = parse_args()

accelerator = Accelerator(
    mixed_precision=args.mixed_precision,
)

# Load scheduler and models
noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="scheduler"
)
noise_scheduler_copy = copy.deepcopy(noise_scheduler)
device = accelerator.device
vae = AutoencoderKL.from_pretrained(
    args.pretrained_model_name_or_path,
    subfolder="vae",
    revision=args.revision,
    variant=args.variant,
)
transformer = FluxTransformer2DModel.from_pretrained(
    args.pretrained_model_name_or_path, subfolder="transformer", revision=args.revision, variant=args.variant
)
transformer.set_frame_embedder()

# We only train the additional adapter LoRA layers
transformer.requires_grad_(False)
vae.requires_grad_(False)

# For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
weight_dtype = torch.float32
if accelerator.mixed_precision == "fp16":
    weight_dtype = torch.float16
elif accelerator.mixed_precision == "bf16":
    weight_dtype = torch.bfloat16

if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
    # due to pytorch#99272, MPS does not yet support bfloat16.
    raise ValueError(
        "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
    )

vae.to(device, dtype=weight_dtype)
transformer.to(device, dtype=weight_dtype)

in_channel = 3
if args.include_focus_depth:
    in_channel += 1
encode_feature_dim = 4096

# from (k,x,y,d) to T5
camera_emb_token_num = 1
camera_projector = torch.nn.Linear(in_channel, encode_feature_dim * camera_emb_token_num) # from (K, x, y) to a positional embedding
camera_projector = camera_projector.to(device=device, dtype=weight_dtype)
camera_projector.requires_grad_(True)
camera_projector_768 = torch.nn.Linear(in_channel, 768).to(device=device, dtype=weight_dtype)

transformer = accelerator.prepare(transformer)
camera_projector = accelerator.prepare(camera_projector)
camera_projector_768 = accelerator.prepare(camera_projector_768)

def get_linear_modules(model):
    linear_modules = []
    
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            linear_modules.append(name)
    
    return linear_modules

target_modules = get_linear_modules(transformer)

lora_rank = 64
transformer_lora_config = LoraConfig(
    r=lora_rank,
    lora_alpha=lora_rank,
    init_lora_weights="gaussian",
    target_modules=target_modules,
)
transformer.add_adapter(transformer_lora_config)

def unwrap_model(model):
    model = accelerator.unwrap_model(model)
    model = model._orig_mod if is_compiled_module(model) else model
    return model

def load_model_hook(models, input_dir):
    transformer_ = None
    camera_projector_ = None
    projector_ = None
    camera_projector_768_ = None

    while len(models) > 0:
        model = models.pop()

        if isinstance(model, type(unwrap_model(transformer))):
            transformer_ = model
        elif model is unwrap_model(camera_projector):
            camera_projector_ = model
        elif model is unwrap_model(camera_projector_768):
            camera_projector_768_ = model
        else:
            raise ValueError(f"unexpected save model: {model.__class__}")

    # load projector weights
    if projector_ is not None:
        projector_state_dict = safetensors.torch.load_file(
            os.path.join(input_dir, "projector.safetensors"),
        )
        projector_.load_state_dict(projector_state_dict)
        print(f"Projector loaded from {os.path.join(input_dir, 'projector.safetensors')}")
    if camera_projector_ is not None:
        camera_projector_state_dict = safetensors.torch.load_file(
            os.path.join(input_dir, "camera_projector.safetensors"),
        )
        camera_projector_.load_state_dict(camera_projector_state_dict)
        print(f"Camera projector loaded from {os.path.join(input_dir, 'camera_projector.safetensors')}")
    if camera_projector_768_ is not None:
        camera_projector_768_state_dict = safetensors.torch.load_file(
            os.path.join(input_dir, "camera_projector_768.safetensors"),
        )
        camera_projector_768_.load_state_dict(camera_projector_768_state_dict)
        print(f"Camera projector 768 loaded from {os.path.join(input_dir, 'camera_projector_768.safetensors')}")
    
    # load lora weights
    lora_state_dict = FluxPipeline.lora_state_dict(input_dir)

    transformer_state_dict = {
        f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")
    }
    transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
    incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")
    if incompatible_keys is not None:
        # check only for unexpected keys
        unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
        if unexpected_keys:
            logger.warning(
                f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
                f" {unexpected_keys}. "
            )

    # Make sure the trainable params are in float32. This is again needed since the base models
    # are in `weight_dtype`. More details:
    # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
    if args.mixed_precision == "fp16":
        models = [transformer_]
        if camera_projector_ is not None:
            models.append(camera_projector_)
        if projector_ is not None:
            models.append(projector_)
        # only upcast trainable parameters (LoRA) into fp32
        cast_training_params(models)

accelerator.register_load_state_pre_hook(load_model_hook)

accelerator.print(f"Resuming from checkpoint {args.resume_from_checkpoint}")
accelerator.load_state(args.resume_from_checkpoint)

transformer = unwrap_model(transformer)
true_cfg_scale = args.cfg_scale
pipeline_args = {
    "vae": vae,
    "transformer": unwrap_model(transformer),
    "scheduler": copy.deepcopy(noise_scheduler),
    "image_lst": None,
    'prompt_embeds': None,
    'pooled_prompt_embeds': None,
    'text_ids': None,
    "height":args.image_size, "width": args.image_size, 
    "num_inference_steps":20, 
    "true_cfg_scale": true_cfg_scale,
}


cnt = 0
input_path = 'asset/input'
image_path = os.path.join(input_path, f'test_refocus_{cnt}.jpg')
depth_path = os.path.join(input_path, f'depth_{cnt}.jpg')

image = Image.open(image_path).convert("RGB").resize((args.image_size, args.image_size))
depth = Image.open(depth_path).convert("L").resize((args.image_size, args.image_size))

image = transforms.ToTensor()(image).to(device, dtype=vae.dtype)
depth = transforms.ToTensor()(depth).to(device, dtype=vae.dtype)


K_scalar = torch.tensor([[args.bokeh_level]]).to(device, dtype=weight_dtype)
focus_coordinates = torch.tensor([[args.focus_point_x, args.focus_point_y]]).to(device, dtype=weight_dtype)
K_scalar = K_scalar / 31.0
projector_input = torch.cat([K_scalar, focus_coordinates], dim=1).to(device, dtype=weight_dtype)


x, y = int(focus_coordinates[:, 0] * args.image_size), int(focus_coordinates[:, 1] * args.image_size)
if args.include_focus_depth:
    d = depth[0, x, y].unsqueeze(0).unsqueeze(0)
    d = d.to(device, dtype=weight_dtype)
    if args.zero_depth:
        d = torch.zeros_like(d).to(device, dtype=weight_dtype)
    projector_input = torch.cat([projector_input, d], dim=1).to(device, dtype=weight_dtype)

if args.include_depth and args.zero_depth:
    depth = torch.zeros_like(depth).to(device, dtype=weight_dtype)

with torch.no_grad():
    camera_emb = camera_projector(projector_input).to(dtype=weight_dtype)
    camera_emb = camera_emb.view(-1, camera_emb_token_num, encode_feature_dim)
    camera_emb_768 = camera_projector_768(projector_input).to(dtype=weight_dtype)
text_ids = torch.zeros(camera_emb.shape[1], 3).to(device=device, dtype=weight_dtype)
if args.include_depth:
    image_lst = [image, depth]
else:
    image_lst = [image]
pipeline_args['image_lst'] = image_lst
pipeline_args['prompt_embeds'] = camera_emb
pipeline_args['pooled_prompt_embeds'] = camera_emb_768
pipeline_args['text_ids'] = text_ids

# generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
generator = None
image_focused = image.clone()
# draw a red rectangle on the image based on the focus coordinates
image_focused[0, x-5:x+5, y-5:y+5] = 1.0
image_focused[1, x-5:x+5, y-5:y+5] = 0.0
image_focused[2, x-5:x+5, y-5:y+5] = 0.0
image_focused = image_focused.cpu().to(torch.float32)
image_focused = transforms.ToPILImage()(image_focused).convert("RGB")
image_focused.save('image_focused.jpg')


import time
for i in tqdm(range(1)):
    start_time = time.time()
    with torch.no_grad():
        pred_image = flux_pipeline_call(**pipeline_args, generator=generator).images[0]
    end_time = time.time()
    print(f"Run {i+1} time: {end_time - start_time:.2f} seconds")
    pred_image = np.array(pred_image)
    cv2.imwrite(f'asset/output/pred_image_{cnt}_{i}.jpg', cv2.cvtColor(pred_image, cv2.COLOR_RGB2BGR))
    print(f'asset/output/pred_image_{cnt}_{i}.jpg saved')