unilm/textdiffuser-2/predict.py

# Prediction interface for Cog ⚙️
# https://github.com/replicate/cog/blob/main/docs/python.md

import copy
import string
import random
from typing import Optional
import torch
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
    StableDiffusionPipeline,
    UNet2DConditionModel,
    DiffusionPipeline,
    LCMScheduler,
)
from tqdm import tqdm
from PIL import Image
from PIL import Image, ImageDraw, ImageFont
from fastchat.model import load_model, get_conversation_template
from transformers import AutoTokenizer, AutoModelForCausalLM
from cog import BasePredictor, Input, Path, BaseModel


alphabet = (
    string.digits
    + string.ascii_lowercase
    + string.ascii_uppercase
    + string.punctuation
    + " "
)  # len(aphabet) = 95
"""alphabet
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ 
"""

font_layout = ImageFont.truetype("./Arial.ttf", 16)


class ModelOutput(BaseModel):
    output_images: list[Path]
    composed_prompt: str


class Predictor(BasePredictor):
    def setup(self) -> None:
        """Load the model into memory to make running multiple predictions efficient"""
        cache_dir = "model_cache"
        local_files_only = True  # set to True if the models are saved in cache_dir

        self.m1_model_path = "JingyeChen22/textdiffuser2_layout_planner"
        self.m1_tokenizer = AutoTokenizer.from_pretrained(
            self.m1_model_path,
            use_fast=False,
            cache_dir=cache_dir,
            local_files_only=local_files_only,
        )
        self.m1_model = AutoModelForCausalLM.from_pretrained(
            self.m1_model_path,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            cache_dir=cache_dir,
            local_files_only=local_files_only,
        ).cuda()

        self.text_encoder = (
            CLIPTextModel.from_pretrained(
                "JingyeChen22/textdiffuser2-full-ft",
                subfolder="text_encoder",
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            .cuda()
            .half()
        )
        self.tokenizer = CLIPTokenizer.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            subfolder="tokenizer",
            cache_dir=cache_dir,
            local_files_only=local_files_only,
        )

        #### additional tokens are introduced, including coordinate tokens and character tokens
        print("***************")
        print(f"tokenizer size: {len(self.tokenizer)}")
        for i in range(520):
            self.tokenizer.add_tokens(["l" + str(i)])  # left
            self.tokenizer.add_tokens(["t" + str(i)])  # top
            self.tokenizer.add_tokens(["r" + str(i)])  # width
            self.tokenizer.add_tokens(["b" + str(i)])  # height
        for c in alphabet:
            self.tokenizer.add_tokens([f"[{c}]"])
        print(f"new tokenizer size: {len(self.tokenizer)}")
        print("***************")

        self.vae = (
            AutoencoderKL.from_pretrained(
                "runwayml/stable-diffusion-v1-5",
                subfolder="vae",
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            .half()
            .cuda()
        )
        self.unet = (
            UNet2DConditionModel.from_pretrained(
                "JingyeChen22/textdiffuser2-full-ft",
                subfolder="unet",
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            .half()
            .cuda()
        )
        self.text_encoder.resize_token_embeddings(len(self.tokenizer))
        self.scheduler = DDPMScheduler.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            subfolder="scheduler",
            cache_dir=cache_dir,
            local_files_only=local_files_only,
        )

        #### load lcm components
        self.pipe = DiffusionPipeline.from_pretrained(
            "lambdalabs/sd-pokemon-diffusers",
            unet=copy.deepcopy(self.unet),
            tokenizer=self.tokenizer,
            text_encoder=copy.deepcopy(self.text_encoder),
            torch_dtype=torch.float16,
            cache_dir=cache_dir,
            local_files_only=local_files_only,
        )
        self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
        self.pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
        self.pipe.to(device="cuda")

    def predict(
        self,
        prompt: str = Input(
            description="Input Prompt. You can let language model automatically identify keywords, or provide them below.",
            default="A beautiful city skyline stamp of Shanghai",
        ),
        keywords: str = Input(
            description="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...).",
            default=None,
        ),
        positive_prompt: str = Input(
            description="(Optional) Positive prompt.",
            default=", digital art, very detailed, fantasy, high definition, cinematic light, dnd, trending on artstation",
        ),
        use_lcm: bool = Input(
            description="Use Latent Consistent Model.", default=False
        ),
        generate_natural_image: bool = Input(
            description="If set to True, the text position and content info will not be incorporated.",
            default=False,
        ),
        num_images: int = Input(
            description="Number of Output images.", default=1, ge=1, le=4
        ),
        num_inference_steps: int = Input(
            description="Number of denoising steps. You may decease the step to 4 when using LCM.",
            ge=1,
            le=50,
            default=20,
        ),
        guidance_scale: float = Input(
            description="Scale for classifier-free guidance. The scale is set to 7.5 by default. When using LCM, guidance_scale is set to 1.",
            ge=1,
            le=20,
            default=7.5,
        ),
        temperature: float = Input(
            description="Control the diversity of layout planner. Higher value indicates more diversity.",
            ge=0.1,
            le=2,
            default=1.4,
        ),
    ) -> ModelOutput:
        """Run a single prediction on the model"""
        if positive_prompt is not None and not len(positive_prompt.strip()) == 0:
            prompt += positive_prompt

        with torch.no_grad():
            user_prompt = prompt

            if generate_natural_image:
                composed_prompt = user_prompt
                prompt = self.tokenizer.encode(user_prompt)
            else:
                if keywords is None or len(keywords.strip()) == 0:
                    template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}"
                else:
                    keywords = keywords.split("/")
                    keywords = [i.strip() for i in keywords]
                    template = f"Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords)}"

                msg = template
                conv = get_conversation_template(self.m1_model_path)
                conv.append_message(conv.roles[0], msg)
                conv.append_message(conv.roles[1], None)
                prompt = conv.get_prompt()
                inputs = self.m1_tokenizer([prompt], return_token_type_ids=False)
                inputs = {k: torch.tensor(v).to("cuda") for k, v in inputs.items()}
                output_ids = self.m1_model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=temperature,
                    repetition_penalty=1.0,
                    max_new_tokens=512,
                )

                if self.m1_model.config.is_encoder_decoder:
                    output_ids = output_ids[0]
                else:
                    output_ids = output_ids[0][len(inputs["input_ids"][0]) :]

                outputs = self.m1_tokenizer.decode(
                    output_ids,
                    skip_special_tokens=True,
                    spaces_between_special_tokens=False,
                )
                print(f"[{conv.roles[0]}]\n{msg}")
                print(f"[{conv.roles[1]}]\n{outputs}")

                ocrs = outputs.split("\n")
                current_ocr = ocrs

                ocr_ids = []
                print("user_prompt", user_prompt)
                print("current_ocr", current_ocr)

                for ocr in current_ocr:
                    ocr = ocr.strip()

                    if len(ocr) == 0 or "###" in ocr or ".com" in ocr:
                        continue

                    items = ocr.split()
                    pred = " ".join(items[:-1])
                    box = items[-1]

                    l, t, r, b = box.split(",")
                    l, t, r, b = int(l), int(t), int(r), int(b)
                    ocr_ids.extend(
                        ["l" + str(l), "t" + str(t), "r" + str(r), "b" + str(b)]
                    )

                    char_list = list(pred)
                    char_list = [f"[{i}]" for i in char_list]
                    ocr_ids.extend(char_list)
                    ocr_ids.append(self.tokenizer.eos_token_id)

                caption_ids = (
                    self.tokenizer(user_prompt, truncation=True, return_tensors="pt")
                    .input_ids[0]
                    .tolist()
                )

                try:
                    ocr_ids = self.tokenizer.encode(ocr_ids)
                    prompt = caption_ids + ocr_ids
                except:
                    prompt = caption_ids

                user_prompt = self.tokenizer.decode(prompt)
                composed_prompt = self.tokenizer.decode(prompt)

            prompt = prompt[:77]
            while len(prompt) < 77:
                prompt.append(self.tokenizer.pad_token_id)

            if not use_lcm:
                prompts_cond = prompt
                prompts_nocond = [self.tokenizer.pad_token_id] * 77

                prompts_cond = [prompts_cond] * num_images
                prompts_nocond = [prompts_nocond] * num_images

                prompts_cond = torch.Tensor(prompts_cond).long().cuda()
                prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()

                scheduler = self.scheduler
                scheduler.set_timesteps(num_inference_steps)
                noise = torch.randn((num_images, 4, 64, 64)).to("cuda").half()
                input = noise

                encoder_hidden_states_cond = self.text_encoder(prompts_cond)[0].half()
                encoder_hidden_states_nocond = self.text_encoder(prompts_nocond)[
                    0
                ].half()

                for t in tqdm(scheduler.timesteps):
                    with torch.no_grad():  # classifier free guidance
                        noise_pred_cond = self.unet(
                            sample=input,
                            timestep=t,
                            encoder_hidden_states=encoder_hidden_states_cond[
                                :num_images
                            ],
                        ).sample  # b, 4, 64, 64
                        noise_pred_uncond = self.unet(
                            sample=input,
                            timestep=t,
                            encoder_hidden_states=encoder_hidden_states_nocond[
                                :num_images
                            ],
                        ).sample  # b, 4, 64, 64
                        noisy_residual = noise_pred_uncond + guidance_scale * (
                            noise_pred_cond - noise_pred_uncond
                        )  # b, 4, 64, 64
                        input = scheduler.step(noisy_residual, t, input).prev_sample
                        del noise_pred_cond
                        del noise_pred_uncond

                        torch.cuda.empty_cache()

                # decode
                input = 1 / self.vae.config.scaling_factor * input
                images = self.vae.decode(input, return_dict=False)[0]
                width, height = 512, 512
                results = []
                new_image = Image.new("RGB", (2 * width, 2 * height))
                for index, image in enumerate(images.cpu().float()):
                    image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
                    image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
                    image = Image.fromarray(
                        (image * 255).round().astype("uint8")
                    ).convert("RGB")
                    results.append(image)
                    row = index // 2
                    col = index % 2
                    new_image.paste(image, (col * width, row * height))
            else:
                generator = torch.Generator(device=self.pipe.device).manual_seed(
                    random.randint(0, 1000)
                )
                results = self.pipe(
                    prompt=user_prompt,
                    generator=generator,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=1,
                    num_images_per_prompt=num_images,
                ).images

        torch.cuda.empty_cache()
        output_paths = []

        for i, sample in enumerate(results):
            output_path = f"/tmp/out-{i}.png"
            sample.save(output_path)
            output_paths.append(Path(output_path))

        return ModelOutput(
            output_images=output_paths,
            composed_prompt=composed_prompt,
        )