predict.py

import os
os.environ["XLA_FLAGS"] = "--xla_gpu_force_compilation_parallelism=1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
import tempfile
import random
from typing import Optional, List
import jax
import jax.numpy as jnp
from flax.jax_utils import replicate
from flax.training.common_utils import shard, shard_prng_key
from functools import partial
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, FlaxCLIPModel
from tqdm.notebook import trange
from cog import BasePredictor, Path, Input, BaseModel

from src.dalle_mini import DalleBart, DalleBartProcessor
from vqgan_jax.modeling_flax_vqgan import VQModel


class ModelOutput(BaseModel):
    clip_score: Optional[float]
    image: Path


class Predictor(BasePredictor):
    def setup(self):
        # Load dalle-mini
        DALLE_MODEL = "checkpoints/dalle_mini_mega-1-fp16"
        DALLE_COMMIT_ID = None
        self.model, self.params = DalleBart.from_pretrained(
            DALLE_MODEL,
            revision=DALLE_COMMIT_ID,
            dtype=jnp.float16,
            _do_init=False,
            ignore_mismatched_sizes=True,
        )
        self.processor = DalleBartProcessor.from_pretrained(
            DALLE_MODEL, revision=DALLE_COMMIT_ID
        )

        # Load VQGAN
        VQGAN_REPO = "checkpoints/vqgan_imagenet_f16_16384"
        VQGAN_COMMIT_ID = "e93a26e7707683d349bf5d5c41c5b0ef69b677a9"

        self.vqgan, self.vqgan_params = VQModel.from_pretrained(
            VQGAN_REPO, revision=VQGAN_COMMIT_ID, _do_init=False
        )

        # Load CLIP
        CLIP_REPO = "checkpoints/openai_clip-vit-base-patch32"
        CLIP_COMMIT_ID = None
        self.clip, self.clip_params = FlaxCLIPModel.from_pretrained(
            CLIP_REPO, revision=CLIP_COMMIT_ID, dtype=jnp.float16, _do_init=False
        )
        self.clip_processor = CLIPProcessor.from_pretrained(
            CLIP_REPO, revision=CLIP_COMMIT_ID
        )

    def predict(
        self,
        prompt: str = Input(
            default="sunset over a lake in the mountains",
            description="Prompt for generating image.",
        ),
        show_clip_score: bool = Input(
            default=False,
            description="CLIP score will be displayed for each generated image if set to True.",
        ),
        n_predictions: int = Input(
            default=8, description="Number of images to generate.", ge=1, le=8
        ),
    ) -> List[ModelOutput]:

        model, params = self.model, self.params
        vqgan, vqgan_params = self.vqgan, self.vqgan_params
        clip, clip_params = self.clip, self.clip_params

        params = replicate(params, device)
        vqgan_params = replicate(vqgan_params, device)
        clip_params = replicate(clip_params, device)
        # exit()

        print("Tokenizing prompt...")
        tokenized_prompt = self.processor([prompt])
        tokenized_prompt = replicate(tokenized_prompt)
        # model inference
        @partial(jax.pmap, axis_name="batch", static_broadcasted_argnums=(3, 4, 5, 6))
        def p_generate(
            tokenized_prompt, key, params, top_k, top_p, temperature, condition_scale
        ):
            return model.generate(
                **tokenized_prompt,
                prng_key=key,
                params=params,
                top_k=top_k,
                top_p=top_p,
                temperature=temperature,
                condition_scale=condition_scale,
            )

        # decode image
        @partial(jax.pmap, axis_name="batch")
        def p_decode(indices, params):
            return vqgan.decode_code(indices, params=params)

        # score images
        @partial(jax.pmap, axis_name="batch")
        def p_clip(inputs, params):
            logits = clip(params=params, **inputs).logits_per_image
            return logits

        # create a random key
        seed = random.randint(0, 2 ** 32 - 1)
        key = jax.random.PRNGKey(seed)

        # We can customize generation parameters
        gen_top_k = None
        gen_top_p = None
        temperature = None
        cond_scale = 3.0

        images = []
        final_output = []

        # generate images
        print("Generating images...")
        for _ in range(n_predictions):
            # get a new key
            key, subkey = jax.random.split(key)
            # generate images
            encoded_images = p_generate(
                tokenized_prompt,
                shard_prng_key(subkey),
                params,
                gen_top_k,
                gen_top_p,
                temperature,
                cond_scale,
            )
            # remove BOS
            encoded_images = encoded_images.sequences[..., 1:]
            # decode images
            decoded_images = p_decode(encoded_images, vqgan_params)
            decoded_images = decoded_images.clip(0.0, 1.0).reshape((-1, 256, 256, 3))

            for decoded_img in decoded_images:
                img = Image.fromarray(np.asarray(decoded_img * 255, dtype=np.uint8))
                images.append(img)

        if not show_clip_score:
            for i, img in enumerate(images):
                out_path = Path(tempfile.mkdtemp()) / f"output_{i}.png"
                img.save(str(out_path))
                final_output.append(ModelOutput(image=out_path))
        else:
            print("Ranking images by CLIP score...")
            # get clip scores
            clip_inputs = self.clip_processor(
                text=[prompt],
                images=images,
                return_tensors="np",
                padding="max_length",
                max_length=77,
                truncation=True,
            ).data
            logits = p_clip(shard(clip_inputs), clip_params)
            logits = logits.squeeze().flatten()

            rank_list = list(logits.argsort()[::-1])

            for i, idx in enumerate(rank_list):
                out_path = Path(tempfile.mkdtemp()) / f"output{idx}.png"
                images[idx].save(str(out_path))
                clip_score = logits[idx]
                final_output.append(ModelOutput(clip_score=clip_score, image=out_path))

        return final_output