run_disenvisioner_w_ip.py

import os
import argparse
import logging
from PIL import Image
from glob import glob
from tqdm import tqdm
from contextlib import nullcontext

import torch
from diffusers import UNet2DConditionModel
from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, CLIPTextModelWithProjection
from diffusers import StableDiffusionPipeline

from disvisioner_modules.disvisioner import DisVisioner
from envisioner_modules.envisioner import Projector, EnVisioner_IP
from envisioner_modules.attention_processor import EVAttnProcessor2_0_W_IP as EVAttnProcessor, AttnProcessor2_0 as AttnProcessor
from envisioner_modules.resampler_masked import Resampler
from utils import seed_all, is_torch2_available, image_grid
assert is_torch2_available()
import global_var


def set_scales(pipe, scale_object, scale_others, scale_ip):
    logging.info(f"==> setting scales: scale_obj {scale_object}, scale_ip {scale_ip}, scale_others {scale_others}")
    for attn_processor in pipe.unet.attn_processors.values():
        if isinstance(attn_processor, EVAttnProcessor):
            attn_processor.scale_object = scale_object
            attn_processor.scale_others = scale_others
            attn_processor.scale_ip = scale_ip

@torch.inference_mode()
def run_inference(args, unet, disv_image_encoder, disv_text_encoder, image_encoder, disvisioner, envisioner, device, dtype):    
    # load SD pipeline
    pipe = StableDiffusionPipeline.from_pretrained(
        args.pretrained_model_name_or_path,
        unet=unet,
        torch_dtype=dtype,
    ).to(device)

    # set scales
    set_scales(pipe, args.scale_object, args.scale_others, args.scale_ip)

    infer_image = args.infer_image
    infer_prompt = args.infer_prompt
    negative_prompt = args.negative_prompt
    print(f"Running for image: {infer_image} with prompt: {infer_prompt}")

    # ------------------ prepare textual embedding ------------------
    class_name = args.class_name
    infer_prompt = infer_prompt.replace("*", class_name)
    print(infer_prompt)

    prompt_embeds_, negative_prompt_embeds_ = pipe.encode_prompt(
        infer_prompt,
        device=device,
        num_images_per_prompt=args.num_samples,
        do_classifier_free_guidance=True,
        negative_prompt=negative_prompt,
    )

    # ------------------ prepare image embedding using DisEnvisioner ------------------
    # read image prompt
    image = Image.open(infer_image)
    image = image.resize((256, 256))
    # get_image_embeds
    clip_image = CLIPImageProcessor()(images=[image], return_tensors="pt").pixel_values
    clip_image = clip_image.to(device, dtype=dtype)
    if torch.backends.mps.is_available():
        autocast_ctx = nullcontext()
    else:
        autocast_ctx = torch.autocast(pipe.device.type)
    with autocast_ctx:
        # disvisioner
        image_features = disv_image_encoder(clip_image.to(device, dtype=dtype), output_hidden_states=True)
        image_embeddings = image_features.last_hidden_state
        image_embeddings = image_embeddings.detach()
        
        class_ids = pipe.tokenizer(
            class_name,
            padding="max_length",
            truncation=True,
            max_length=pipe.tokenizer.model_max_length,
            return_tensors="pt",
        ).input_ids[0].to(device).unsqueeze(0)
        
        class_proj = disv_text_encoder(class_ids.to(device)).text_embeds
        inj_embedding, obj_map, _ = disvisioner(image_embeddings, class_proj, return_attns=True)
        obj_mask =(obj_map > 0.3).int()

        # NOTE
        # currently the mask is passed through global_var for convinience, 
        # thus, the batchsize must be 1!
        assert obj_mask.shape[0] == 1
        global_var.set_value("mask", obj_mask[0].cpu())

        # envisioner
        # get projected embeds
        image_prompt_embeds_object, image_prompt_embeds_others = envisioner.image_proj_model(inj_embedding.float())            
        uncond_image_prompt_embeds_object, uncond_image_prompt_embeds_others = envisioner.image_proj_model(torch.zeros_like(inj_embedding.float()))

        bs_embed, seq_len, _ = image_prompt_embeds_object.shape
        image_prompt_embeds_object = image_prompt_embeds_object.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)
        uncond_image_prompt_embeds_object = uncond_image_prompt_embeds_object.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)
        
        bs_embed, seq_len, _ = image_prompt_embeds_others.shape
        image_prompt_embeds_others = image_prompt_embeds_others.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)
        uncond_image_prompt_embeds_others = uncond_image_prompt_embeds_others.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)

        image_embeddings_ip = image_encoder(clip_image, output_hidden_states=True).hidden_states[-2]
        uncond_image_embeddings_ip = image_encoder(torch.zeros_like(clip_image), output_hidden_states=True).hidden_states[-2]

        image_prompt_embeds_ip = envisioner.image_proj_model_ip(image_embeddings_ip)
        uncond_image_prompt_embeds_ip = envisioner.image_proj_model_ip(uncond_image_embeddings_ip)

        bs_embed, seq_len, _ = image_prompt_embeds_ip.shape
        image_prompt_embeds_ip = image_prompt_embeds_ip.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)
        uncond_image_prompt_embeds_ip = uncond_image_prompt_embeds_ip.repeat(1, args.num_samples, 1).view(bs_embed * args.num_samples, seq_len, -1)

        
        # ------------------ prepare textual and image embeddings ------------------
        prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds_object, image_prompt_embeds_others, image_prompt_embeds_ip], dim=1)
        negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds_object, uncond_image_prompt_embeds_others, uncond_image_prompt_embeds_ip], dim=1)

        # gen
        generator = torch.Generator(device).manual_seed(args.seed) if args.seed is not None else None
        images = pipe(
            prompt_embeds=prompt_embeds,
            height=args.resolution,
            width=args.resolution,
            negative_prompt_embeds=negative_prompt_embeds,
            guidance_scale=args.guidance_scale,
            num_inference_steps=args.num_inference_steps,
            generator=generator
        ).images

    # save
    gen_images = [image.resize((args.resolution, args.resolution))] + images
    grid = image_grid(gen_images, len(gen_images)//(1+args.num_samples), 1+args.num_samples)
    # val_file_name = '.'.join(os.path.basename(infer_image).split(".")[:-1])
    grid.save(os.path.join(args.output_dir, f'sobj{args.scale_object}_sip{args.scale_ip}_soth{args.scale_others}_[{infer_prompt}]_seed{args.seed}.png'))

    
def parse_args():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        default=None,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--half_precision",
        action="store_true"
    )
    parser.add_argument(
        "--disvisioner_path",
        type=str,
        default=None,
        help="Path to pretrained disvisioner.",
        required=True
    )
    parser.add_argument(
        "--token_num",
        type=int,
        default=1,
        help="number of tokens for object"
    )
    parser.add_argument(
        "--pretrained_CLIP",
        type=str,
        default=None,
        help="Path to pretrained disvisioner encoders.",
        required=True
    )
    parser.add_argument(
        "--ip_image_encoder_path",
        type=str,
    )
    parser.add_argument(
        "--scale_object",
        type=float,
        default=0.8,
    )
    parser.add_argument(
        "--scale_others",
        type=float,
        default=0.0,
    )
    parser.add_argument(
        "--scale_ip",
        type=float,
        default=0.0,
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="disenvisioner",
        help="The output directory where the model predictions will be written.",
    )
    parser.add_argument(
        "--resolution",
        type=int,
        default=512,
        help=(
            "The resolution for input images"
        ),
    )
    parser.add_argument(
        "--infer_image",
        type=str,
        required=True
    )
    parser.add_argument(
        "--infer_prompt",
        type=str,
        required=True
    )
    parser.add_argument(
        "--negative_prompt",
        type=str,
        default="monochrome, lowres, bad anatomy, worst quality, low quality"
    )
    parser.add_argument(
        "--class_name",
        type=str,
        required=True
    )
    parser.add_argument(
        "--num_samples",
        type=int,
        default=4
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
        default=7.5
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=50
    )
    parser.add_argument(
        "--envisioner_path",
        type=str,
        default=None,
        help="Path to pretrained envisioner.",
    )
    parser.add_argument(
        "--object_factor",
        type=int,
        default=1,
        help="factor that determines the number of object tokens. "
    )
    parser.add_argument(
        "--others_factor",
        type=int,
        default=1,
        help="factor that determines the number of other component tokens."
    ) 
    args = parser.parse_args()

    return args
    
def main():
    logging.basicConfig(level=logging.INFO)
    logging.info(f"Run inference...")
    
    args = parse_args()
    global_var._init()
    # ------------------ Preparation ------------------
    os.makedirs(args.output_dir, exist_ok=True)
    logging.info(f"Output dir = {args.output_dir}")

    if args.seed is not None:
        seed_all(args.seed)

    # half_precision
    if args.half_precision:
        dtype = torch.float16
        logging.info(f"Running with half precision ({dtype}).")
    else:
        dtype = torch.float32

    # -------------------- Device --------------------
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
        logging.warning("CUDA is not available. Running on CPU will be slow.")
    logging.info(f"Device = {device}")

    # -------------------- Models --------------------
    # ------------------ 1. Load Unet model and image encoder for ip-adapter ------------------
    unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.ip_image_encoder_path)
    
    # ------------------ 2. Load models for disvisioner ------------------
    disv_text_encoder = CLIPTextModelWithProjection.from_pretrained(args.pretrained_CLIP)
    disv_image_encoder = CLIPVisionModelWithProjection.from_pretrained(args.pretrained_CLIP)
    # define disvisioner
    disvisioner = DisVisioner(
        image_hidden_size=disv_image_encoder.vision_model.config.hidden_size, 
        text_hidden_size=disv_text_encoder.text_model.config.hidden_size,
        output_dim=disv_text_encoder.text_model.config.hidden_size,
        token_num=args.token_num, num_refine=2
        )
    logging.info(f"Load Disvisioner from {args.disvisioner_path}")
    disvisioner.load_state_dict(torch.load(args.disvisioner_path, map_location='cpu'), strict=True)

    # ------------------ 3. Load envisioner ------------------
    # 1.1 define projector
    image_projector = Projector(
        cross_attention_dim=unet.config.cross_attention_dim, # output dim
        input_embedding_dim=disv_text_encoder.text_model.config.hidden_size, # input dim
        clip_extra_context_tokens=4,
    )
    # 1.2 define image projection model for ip-adapter
    image_proj_model_ip = Resampler(
        dim=unet.config.cross_attention_dim,
        depth=4,
        dim_head=64,
        heads=12,
        num_queries=16,
        embedding_dim=image_encoder.config.hidden_size,
        output_dim=unet.config.cross_attention_dim,
        ff_mult=4
    )
    # 2. Define additional CA modules (referred to IP-Adapter)
    attn_procs = {}
    for name in unet.attn_processors.keys():
        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
        if name.startswith("mid_block"):
            hidden_size = unet.config.block_out_channels[-1]
        elif name.startswith("up_blocks"):
            block_id = int(name[len("up_blocks.")])
            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
        elif name.startswith("down_blocks"):
            block_id = int(name[len("down_blocks.")])
            hidden_size = unet.config.block_out_channels[block_id]
        
        if cross_attention_dim is None: # attn1
            attn_procs[name] = AttnProcessor()
        else: # attn2
            attn_procs[name] = EVAttnProcessor(hidden_size=hidden_size, 
                                               cross_attention_dim=cross_attention_dim, 
                                               num_tokens_object=args.object_factor*4, # number of object tokens
                                               num_tokens_others=args.others_factor*4, # number of other tokens
                                               num_tokens_ip=16
                                               )
    unet.set_attn_processor(attn_procs)
    adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())

    # 3. Define envisioner (Projector + added adapter_modules)
    envisioner = EnVisioner_IP(image_projector, image_proj_model_ip, adapter_modules)
    
    logging.info(f"Load EnVisioner from {args.envisioner_path}")
    envisioner.load_state_dict(torch.load(args.envisioner_path),strict=True)

    unet.to(device, dtype=dtype)
    disv_image_encoder.to(device, dtype=dtype)
    disv_text_encoder.to(device, dtype=dtype)
    disvisioner.to(device, dtype=dtype)
    envisioner.to(device, dtype=dtype)
    image_encoder.to(device,  dtype=dtype)

    run_inference(args, unet, disv_image_encoder, disv_text_encoder, image_encoder, disvisioner, envisioner, device, dtype)


if __name__ == "__main__":
    main()