infer.py

import os
import argparse
import numpy as np
import torch

from PIL import Image
from schedulers.lcm_single_step_scheduler import LCMSingleStepScheduler

from diffusers import (
    DDPMScheduler,
    StableDiffusionXLPipeline
)

from transformers import (
    CLIPImageProcessor, CLIPVisionModelWithProjection,
    AutoImageProcessor, AutoModel
)

from module.ip_adapter.utils import init_adapter_in_unet
from module.ip_adapter.resampler import Resampler
from pipelines.sdxl_instantir import InstantIRPipeline, PREVIEWER_LORA_MODULES, LCM_LORA_MODULES


def name_unet_submodules(unet):
    def recursive_find_module(name, module, end=False):
        if end:
            for sub_name, sub_module in module.named_children():
                sub_module.full_name = f"{name}.{sub_name}"
            return
        if not "up_blocks" in name and not "down_blocks" in name and not "mid_block" in name: return
        elif "resnets" in name: return
        for sub_name, sub_module in module.named_children():
            end = True if sub_name == "transformer_blocks" else False
            recursive_find_module(f"{name}.{sub_name}", sub_module, end)

    for name, module in unet.named_children():
        recursive_find_module(name, module)


def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
               pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):

    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        # ratio = min_side / min(h, w)
        # w, h = round(ratio*w), round(ratio*h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
        input_image = Image.fromarray(res)
    return input_image


def tensor_to_pil(images):
    """
    Convert image tensor or a batch of image tensors to PIL image(s).
    """
    images = images.clamp(0, 1)
    images_np = images.detach().cpu().numpy()
    if images_np.ndim == 4:
        images_np = np.transpose(images_np, (0, 2, 3, 1))
    elif images_np.ndim == 3:
        images_np = np.transpose(images_np, (1, 2, 0))
        images_np = images_np[None, ...]
    images_np = (images_np * 255).round().astype("uint8")
    if images_np.shape[-1] == 1:
        # special case for grayscale (single channel) images
        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images_np]
    else:
        pil_images = [Image.fromarray(image[:, :, :3]) for image in images_np]

    return pil_images


def calc_mean_std(feat, eps=1e-5):
	"""Calculate mean and std for adaptive_instance_normalization.
	Args:
		feat (Tensor): 4D tensor.
		eps (float): A small value added to the variance to avoid
			divide-by-zero. Default: 1e-5.
	"""
	size = feat.size()
	assert len(size) == 4, 'The input feature should be 4D tensor.'
	b, c = size[:2]
	feat_var = feat.view(b, c, -1).var(dim=2) + eps
	feat_std = feat_var.sqrt().view(b, c, 1, 1)
	feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
	return feat_mean, feat_std


def adaptive_instance_normalization(content_feat, style_feat):
    size = content_feat.size()
    style_mean, style_std = calc_mean_std(style_feat)
    content_mean, content_std = calc_mean_std(content_feat)
    normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
    return normalized_feat * style_std.expand(size) + style_mean.expand(size)


def main(args, device):

    # image encoder and feature extractor.
    if args.use_clip_encoder:
        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
            args.vision_encoder_path,
            subfolder="image_encoder",
        )
        image_processor = CLIPImageProcessor()
    else:
        image_encoder = AutoModel.from_pretrained(args.vision_encoder_path)
        image_processor = AutoImageProcessor.from_pretrained(args.vision_encoder_path)
    image_encoder.to(torch.float16)

    # Base models.
    pipe = StableDiffusionXLPipeline.from_pretrained(
        args.sdxl_path,
        torch_dtype=torch.float16,
        revision=args.revision,
        variant=args.variant
    )

    # InstantIR pipeline
    pipe = InstantIRPipeline(
            pipe.vae, pipe.text_encoder, pipe.text_encoder_2, pipe.tokenizer, pipe.tokenizer_2,
            pipe.unet, pipe.scheduler, feature_extractor=image_processor, image_encoder=image_encoder,
    ).to(device)
    unet = pipe.unet

    # Image prompt projector.
    print("Loading LQ-Adapter...")
    image_proj_model = Resampler(
        embedding_dim=image_encoder.config.hidden_size,
        output_dim=unet.config.cross_attention_dim,
    )
    adapter_path = args.adapter_model_path if args.adapter_model_path is not None else os.path.join(args.instantir_path, 'adapter.pt')
    init_adapter_in_unet(
        unet,
        image_proj_model,
        adapter_path,
    )

    # Prepare previewer
    previewer_lora_path = args.previewer_lora_path if args.previewer_lora_path is not None else args.instantir_path
    if previewer_lora_path is not None:
        lora_alpha = pipe.prepare_previewers(previewer_lora_path)
        print(f"use lora alpha {lora_alpha}")
    unet.to(device, dtype=torch.float16)
    pipe.scheduler = DDPMScheduler.from_pretrained(args.sdxl_path, subfolder="scheduler")
    lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)

    # Load weights.
    print("Loading checkpoint...")
    pretrained_state_dict = torch.load(os.path.join(args.instantir_path, "aggregator.pt"), map_location="cpu")
    pipe.aggregator.load_state_dict(pretrained_state_dict, strict=True)
    pipe.aggregator.to(device, dtype=torch.float16)

    #################### Restoration ####################

    post_fix = f"_{args.post_fix}" if args.post_fix else ""
    post_fix = args.instantir_path.split("/")[-2]+f"{post_fix}"
    os.makedirs(f"{args.out_path}/{post_fix}", exist_ok=True)

    processed_imgs = os.listdir(os.path.join(args.out_path, post_fix))
    lq_files = []
    lq_batch = []
    for file in os.listdir(args.test_path):
        if file in processed_imgs:
            print(f"Skip {file}")
            continue
        lq_batch.append(f"{file}")
        if len(lq_batch) == args.batch_size:
            lq_files.append(lq_batch)
            lq_batch = []

    if len(lq_batch) > 0:
        lq_files.append(lq_batch)

    for lq_batch in lq_files:
        generator = torch.Generator(device=device).manual_seed(args.seed)
        pil_lqs = [Image.open(os.path.join(args.test_path, file)) for file in lq_batch]
        if args.width is None or args.height is None:
            lq = [resize_img(pil_lq.convert("RGB"), size=None) for pil_lq in pil_lqs]
        else:
            lq = [resize_img(pil_lq.convert("RGB"), size=(args.width, args.height)) for pil_lq in pil_lqs]
        timesteps = None
        if args.denoising_start < 1000:
            timesteps = [
                i * (args.denoising_start//args.num_inference_steps) + pipe.scheduler.config.steps_offset for i in range(0, args.num_inference_steps)
            ]
            timesteps = timesteps[::-1]
            pipe.scheduler.set_timesteps(args.num_inference_steps, device)
            timesteps = pipe.scheduler.timesteps
        prompt = args.prompt
        if not isinstance(prompt, list):
            prompt = [prompt]
        prompt = prompt*len(lq)
        neg_prompt = args.neg_prompt
        if not isinstance(neg_prompt, list):
            neg_prompt = [neg_prompt]
        neg_prompt = neg_prompt*len(lq)
        image = pipe(
            prompt=prompt,
            image=lq,
            ip_adapter_image=[lq],
            num_inference_steps=args.num_inference_steps,
            generator=generator,
            timesteps=timesteps,
            negative_prompt=neg_prompt,
            guidance_scale=args.cfg,
            previewer_scheduler=lcm_scheduler,
            return_dict=False,
        )[0]

        if args.save_preview_row:
            for i, lcm_image in enumerate(image[1]):
                lcm_image.save(f"./lcm/{i}.png")
        for i, rec_image in enumerate(image):
            rec_image.save(f"{args.out_path}/{post_fix}/{lq_batch[i]}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="InstantIR pipeline")
    parser.add_argument(
        "--sdxl_path",
        type=str,
        default=None,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--previewer_lora_path",
        type=str,
        default=None,
        help="Path to LCM lora or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--pretrained_vae_model_name_or_path",
        type=str,
        default=None,
        help="Path to an improved VAE to stabilize training. For more details check out: https://github.com/huggingface/diffusers/pull/4038.",
    )
    parser.add_argument(
        "--instantir_path",
        type=str,
        default=None,
        required=True,
        help="Path to pretrained instantir model.",
    )
    parser.add_argument(
        "--vision_encoder_path",
        type=str,
        default='/share/huangrenyuan/model_zoo/vis_backbone/dinov2_large',
        help="Path to image encoder for IP-Adapters or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--adapter_model_path",
        type=str,
        default=None,
        help="Path to IP-Adapter models or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--adapter_tokens",
        type=int,
        default=64,
        help="Number of tokens to use in IP-adapter cross attention mechanism.",
    )
    parser.add_argument(
        "--use_clip_encoder",
        action="store_true",
        help="Whether or not to use DINO as image encoder, else CLIP encoder.",
    )
    parser.add_argument(
        "--denoising_start",
        type=int,
        default=1000,
        help="Diffusion start timestep."
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=30,
        help="Diffusion steps."
    )
    parser.add_argument(
        "--resolution",
        type=int,
        default=1024,
        help="Number of tokens to use in IP-adapter cross attention mechanism.",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=6,
        help="Test batch size."
    )
    parser.add_argument(
        "--width",
        type=int,
        default=None,
        help="Output image width."
    )
    parser.add_argument(
        "--height",
        type=int,
        default=None,
        help="Output image height."
    )
    parser.add_argument(
        "--cfg",
        type=float,
        default=7.0,
        help="Scale of Classifier-Free-Guidance (CFG).",
    )
    parser.add_argument(
        "--post_fix",
        type=str,
        default=None,
        help="Subfolder name for restoration output under the output directory.",
    )
    parser.add_argument(
        "--variant",
        type=str,
        default='fp16',
        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
    )
    parser.add_argument(
        "--revision",
        type=str,
        default=None,
        required=False,
        help="Revision of pretrained model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--save_preview_row",
        action="store_true",
        help="Whether or not to save the intermediate lcm outputs.",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default='',
        nargs="+",
        help=(
            "A set of prompts for creative restoration. Provide either a matching number of test images,"
            " or a single prompt to be used with all inputs."
        ),
    )
    parser.add_argument(
        "--neg_prompt",
        type=str,
        default='',
        nargs="+",
        help=(
            "A set of negative prompts for creative restoration. Provide either a matching number of test images,"
            " or a single negative prompt to be used with all inputs."
        ),
    )
    parser.add_argument(
        "--test_path",
        type=str,
        default=None,
        required=True,
        help="Test directory.",
    )
    parser.add_argument(
        "--out_path",
        type=str,
        default="./output",
        help="Output directory.",
    )
    parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
    args = parser.parse_args()
    args.height = args.height or args.width
    args.width = args.width or args.height
    if args.width % 64 != 0 or args.height % 64 != 0:
        raise ValueError("Image resolution must be divisible by 64.")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    main(args, device)