gen_story.py

from models.transforms import get_object_transforms
from models.data import EvalStoryDataset, EvalPororoStoryDataset
from models.model import StoryModel
from diffusers import StableDiffusionPipeline
from transformers import CLIPTokenizer
from accelerate.utils import set_seed
from models.utils import parse_args
from accelerate import Accelerator
from pathlib import Path
from PIL import Image
import numpy as np
import torch
from torchvision import transforms
import os
from tqdm.auto import tqdm
from models.pipeline import (
    stable_diffusion_call_with_references_delayed_conditioning,
)
import types
import os
import random
import re

from gill import models
from gill import utils

from torchvision import transforms


@torch.no_grad()
def main():
    args = parse_args()
    accelerator = Accelerator(
        mixed_precision=args.mixed_precision,
    )

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    pipe = StableDiffusionPipeline.from_pretrained(
        args.pretrained_model_name_or_path, torch_dtype=weight_dtype
    ).to(accelerator.device)
    pipe.set_progress_bar_config(disable=True)
    pipe.safety_checker = None
    pipe.requires_safety_checker = False

    model = StoryModel.from_pretrained(args)
    model.eval()

    ckpt_name = "pytorch_model.bin"

    model.load_state_dict(
        torch.load(Path(args.finetuned_model_path) / ckpt_name, map_location="cpu"), strict=False
    )

    model = model.to(device=accelerator.device, dtype=weight_dtype)

    model_dir = args.gill_ckpt
    mm_llm = models.load_gill(model_dir, device=accelerator.device)
    mm_llm.eval()
    g_cuda = torch.Generator(device=accelerator.device).manual_seed(1337)

    pipe.unet = model.unet

    if args.enable_xformers_memory_efficient_attention:
        pipe.unet.enable_xformers_memory_efficient_attention()

    pipe.text_encoder = model.text_encoder
    pipe.image_encoder = model.image_encoder

    pipe.postfuse_module = model.postfuse_module

    pipe.inference = types.MethodType(
        stable_diffusion_call_with_references_delayed_conditioning, pipe
    )

    del model

    # Set up the dataset
    tokenizer = CLIPTokenizer.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="tokenizer",
        revision=args.revision,
    )

    object_transforms = get_object_transforms(args)

    if args.dataset == 'flintstones':
        demo_dataset = EvalStoryDataset(
            tokenizer=tokenizer,
            object_transforms=object_transforms,
            device=accelerator.device,
            max_num_objects=args.max_num_objects,
            root=args.dataset_name,
            ref_image=args.ref_image,
            story_len=args.story_len
        )
    else:
        demo_dataset = EvalPororoStoryDataset(
            tokenizer=tokenizer,
            object_transforms=object_transforms,
            device=accelerator.device,
            max_num_objects=args.max_num_objects,
            root=args.dataset_name,
            ref_image=args.ref_image
        )
    os.makedirs(args.output_dir, exist_ok=True)

    image_id = 's_01_e_25_shot_004015_004089'
    
    batchs = demo_dataset.prepare_data_batch(image_id)
    gen_images = []
    gen_captions = batchs[0]['captions']
    for idx, batch in enumerate(batchs):
        prompt_llm = []
        if idx == 0:
            prompt_llm.append('Caption: ' + batch['captions'][idx] + ' Image: ')
        else:
            for i in range(idx + 1):
                if i == idx:
                    prompt_llm.append(' Caption: ' + batch['captions'][i] + ' Image: ')
                else:
                    prompt_llm.append(' Image: <img>')
                    prompt_llm.append(gen_images[i])
                    prompt_llm.append('</img> Caption: ' + batch['captions'][i])
        input_ids = batch["input_ids"].to(accelerator.device)
        image_id = batch['image_id']
        image_token_mask = batch["image_token_mask"].to(accelerator.device)
        all_object_pixel_values = (
            batch["object_pixel_values"].unsqueeze(0).to(accelerator.device)
        )
        num_objects = batch["num_objects"].unsqueeze(0).to(accelerator.device)

        all_object_pixel_values = all_object_pixel_values.to(
            dtype=weight_dtype, device=accelerator.device
        )

        object_pixel_values = all_object_pixel_values  # [:, 0, :, :, :]

        if pipe.image_encoder is not None:
            object_embeds = pipe.image_encoder(object_pixel_values)
        else:
            object_embeds = None

        if idx == 0:
            encoder_hidden_states = pipe.text_encoder(input_ids)[0]
            encoder_hidden_states = pipe.postfuse_module(
                encoder_hidden_states,
                object_embeds,
                image_token_mask,
                num_objects,
            )
            start_merge_step = args.start_merge_step
        else:
            encoder_hidden_states, caption = mm_llm.generate_for_images_emb(prompt_llm, num_words=36, min_word_tokens=30)
            encoder_hidden_states = encoder_hidden_states.half()
            if batch['ref_flag']:
                start_merge_step = 0
            else:
                start_merge_step = args.start_merge_step

        encoder_hidden_states_text_only = pipe._encode_prompt(
            batch['prompt_text_only'],
            accelerator.device,
            args.num_images_per_prompt,
            do_classifier_free_guidance=False,
        )

        cross_attention_kwargs = {}

        images = pipe.inference(
            prompt_embeds=encoder_hidden_states,
            num_inference_steps=args.inference_steps,
            height=args.generate_height,
            width=args.generate_width,
            guidance_scale=args.guidance_scale,
            num_images_per_prompt=args.num_images_per_prompt,
            cross_attention_kwargs=cross_attention_kwargs,
            prompt_embeds_text_only=encoder_hidden_states_text_only,
            start_merge_step=start_merge_step,
        ).images

        gen_images.append(images[0])

        cap_img = utils.create_image_of_text(batch['prompt_text_only'],
                                                width=args.generate_width, nrows=2, color=(255, 255, 0))
        image = transforms.ToTensor()(images[0])
        image = torch.cat([image, cap_img], dim=1)
        pil_image = transforms.ToPILImage()(image)
        pil_image.save(os.path.join(args.output_dir, f"{image_id}.png"))
    
    max_length = 10
    gen_caption_flag = True
    while len(gen_images) < max_length:
        prompt_llm = []
        if gen_caption_flag:
            for i in range(len(gen_captions) - 4, len(gen_captions)):
                prompt_llm.append('Caption: ' + gen_captions[i])
            prompt_llm.append(f' Generate the next frame description of cartoon Flintstones based on the previous captions. The new generation should be different from the previous captions.')
        else:
            for i in range(len(gen_images) - 4, len(gen_images)):
                prompt_llm.append(' Image: <img>')
                prompt_llm.append(gen_images[i])
                prompt_llm.append('</img> Caption: ' + gen_captions[i])
            prompt_llm.append(f' Caption: {gen_captions[-1]} Image: ')
        
        if gen_caption_flag:
            encoder_hidden_states, caption = mm_llm.generate_for_images_emb(prompt_llm, num_words=36, min_word_tokens=30)
            caption = caption.replace('\n', '')
            index = caption.find("[IMG")
            if index != -1:
                caption = caption[:index]
            index = caption.find(".")
            if index != -1:
                caption = caption[:index+1]
            index = caption.find(":")
            if index != -1:
                caption = caption[index+1:]
            caption = caption.strip()
            gen_captions.append(caption)
            gen_caption_flag = False
        else:
            encoder_hidden_states, caption = mm_llm.generate_for_images_emb(prompt_llm, num_words=2)
            encoder_hidden_states = encoder_hidden_states.half()
            images = pipe.inference(
                prompt_embeds=encoder_hidden_states,
                num_inference_steps=args.inference_steps,
                height=args.generate_height,
                width=args.generate_width,
                guidance_scale=args.guidance_scale,
                num_images_per_prompt=args.num_images_per_prompt,
                cross_attention_kwargs=cross_attention_kwargs,
                prompt_embeds_text_only=encoder_hidden_states,
                start_merge_step=0,
            ).images
            gen_images.append(images[0])
            gen_caption_flag = True
            cap_img = utils.create_image_of_text(gen_captions[-1], width=args.generate_width, nrows=2, color=(255, 255, 0))
            image = transforms.ToTensor()(images[0])
            image = torch.cat([image, cap_img], dim=1)
            pil_image = transforms.ToPILImage()(image)
            pil_image.save(os.path.join(args.output_dir, f"{len(gen_images)}.png"))        
        

if __name__ == "__main__":
    main()