init

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jinheng Xie
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4081eb457c27efd213e1df5d71e6075fdd1a969a7ac408bf8e7968345250d360
-size 227
+oid sha256:df8da49db1cd3db14e34d015e398bd3d5ab51c3988fd98976405701ce1838ef5
+size 224

app.py

@@ -0,0 +1,526 @@
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "False"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+import numpy as np
+import gradio as gr
+import spaces
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+
+from prompting_utils import UniversalPrompting, create_attention_mask_predict_next, create_attention_mask_for_mmu
+from training_utils import image_transform
+from models import Showo, MAGVITv2, get_mask_chedule
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+config = OmegaConf.load("configs/showo_demo.yaml")
+tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                   special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>",
+                                                   "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                   ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+vq_model = MAGVITv2()
+vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+vq_model.requires_grad_(False)
+vq_model.eval()
+
+model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+model.eval()
+mask_token_id = model.config.mask_token_id
+
+
+@spaces.GPU
+def text_to_image_generation(input_text, guidance_scale=1.75, generation_timesteps=18):
+    prompts = [input_text]
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    image_tokens = torch.ones((len(prompts), config.model.showo.num_vq_tokens),
+                              dtype=torch.long, device=device) * mask_token_id
+
+    input_ids, _ = uni_prompting((prompts, image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompts), image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+@spaces.GPU
+def text_guided_inpainting(input_text, inpainting_image, inpainting_mask, guidance_scale=1.75, generation_timesteps=16):
+    prompt = [input_text]
+
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    inpainting_image = image_transform(inpainting_image, resolution=config.dataset.params.resolution).to(device)
+    inpainting_mask = image_transform(inpainting_mask, resolution=config.dataset.params.resolution, normalize=False)
+
+    inpainting_image = inpainting_image.unsqueeze(0).repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask = inpainting_mask.unsqueeze(0).to(device)
+    inpainting_mask = F.interpolate(inpainting_mask, size=config.dataset.params.resolution // 16, mode='bicubic')
+    inpainting_mask = inpainting_mask.repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask[inpainting_mask < 0.5] = 0
+    inpainting_mask[inpainting_mask >= 0.5] = 1
+
+    inpainting_mask = inpainting_mask.reshape(config.training.batch_size, -1)
+    inpainting_mask = inpainting_mask.to(torch.bool)
+
+    inpainting_image_tokens = vq_model.get_code(inpainting_image) + len(uni_prompting.text_tokenizer)
+    inpainting_image_tokens[inpainting_mask] = mask_token_id
+
+    input_ids, _ = uni_prompting((prompt, inpainting_image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompt), inpainting_image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+@spaces.GPU
+def text_guided_extrapolation(input_img, input_text, left_ext, right_ext, guidance_scale=1.75, generation_timesteps=16):
+    config.offset = 0
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    extra_direction = ['right'] * int(right_ext) + ['left'] * int(left_ext)
+    prompt = [input_text] * len(extra_direction)
+    W = config.dataset.params.resolution // 16
+    for id, (prt, direction) in enumerate(zip(prompt, extra_direction)):
+        prt = [prt] * config.training.batch_size
+        if id == 0:
+            # extrapolation_image = Image.open(config.image_path).convert("RGB")
+            extrapolation_image = input_img
+            extrapolation_image = image_transform(extrapolation_image,
+                                                  resolution=config.dataset.params.resolution).to(device)
+
+            B, _, _ = extrapolation_image.shape
+            extrapolation_image = extrapolation_image.unsqueeze(0)
+            extrapolation_image_tokens = vq_model.get_code(extrapolation_image) + len(uni_prompting.text_tokenizer)
+            extrapolation_image_tokens = extrapolation_image_tokens.reshape(1,
+                                                                            config.dataset.params.resolution // 16,
+                                                                            config.dataset.params.resolution // 16)
+            extrapolation_image_tokens = extrapolation_image_tokens.repeat(config.training.batch_size, 1, 1)
+        else:
+
+            extrapolation_image_tokens = gen_token_ids + len(uni_prompting.text_tokenizer)
+
+        image_left_part = extrapolation_image_tokens[:, :, :-(W // 2 - config.offset)] - len(
+            uni_prompting.text_tokenizer)
+        image_right_part = extrapolation_image_tokens[:, :, W // 2 - config.offset:] - len(uni_prompting.text_tokenizer)
+        image_up_part = extrapolation_image_tokens[:, :-(W // 2 - config.offset), :] - len(uni_prompting.text_tokenizer)
+        image_down_part = extrapolation_image_tokens[:, W // 2 - config.offset:, :] - len(uni_prompting.text_tokenizer)
+
+        if direction in ['left', 'right']:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset),
+                                             dtype=torch.int64, device=device) + mask_token_id
+        else:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset,
+                                              config.dataset.params.resolution // 16),
+                                             dtype=torch.int64, device=device) + mask_token_id
+
+        if direction == 'left':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :, :W // 2 - config.offset]], dim=-1)
+        elif direction == 'right':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, :, -(W // 2 - config.offset):], extrapolation_mask], dim=-1)
+        elif direction == 'up':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :W // 2 - config.offset, :]], dim=-2)
+        else:
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, -(W // 2 - config.offset):, :], extrapolation_mask], dim=-2)
+
+        extrapolation_image_tokens = extrapolation_image_tokens.reshape(config.training.batch_size, -1)
+
+        input_ids, _ = uni_prompting((prt, extrapolation_image_tokens), 't2i_gen')
+
+        if config.training.guidance_scale > 0:
+            uncond_input_ids, _ = uni_prompting(([''] * len(prt), extrapolation_image_tokens), 't2i_gen')
+            attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+        else:
+            attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+            uncond_input_ids = None
+
+        if config.get("mask_schedule", None) is not None:
+            schedule = config.mask_schedule.schedule
+            args = config.mask_schedule.get("params", {})
+            mask_schedule = get_mask_chedule(schedule, **args)
+        else:
+            mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+        with torch.no_grad():
+            gen_token_ids = model.t2i_generate(
+                input_ids=input_ids,
+                uncond_input_ids=uncond_input_ids,
+                attention_mask=attention_mask,
+                guidance_scale=config.training.guidance_scale,
+                temperature=config.training.get("generation_temperature", 1.0),
+                timesteps=config.training.generation_timesteps,
+                noise_schedule=mask_schedule,
+                noise_type=config.training.get("noise_type", "mask"),
+                seq_len=config.model.showo.num_vq_tokens,
+                uni_prompting=uni_prompting,
+                config=config,
+            )
+
+        gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+        gen_token_ids = gen_token_ids.reshape(config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16)
+        if direction == 'left':
+            gen_token_ids = torch.cat([gen_token_ids, image_right_part], dim=-1)
+        elif direction == 'right':
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-1)
+        elif direction == 'up':
+            gen_token_ids = torch.cat([gen_token_ids, image_down_part], dim=-2)
+        else:
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-2)
+
+    _, h, w = gen_token_ids.shape
+    gen_token_ids = gen_token_ids.reshape(config.training.batch_size, -1)
+    images = vq_model.decode_code(gen_token_ids, shape=(h, w))
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+@spaces.GPU
+def multimodal_understanding(input_img, input_text, chat_history):
+    top_k = 1  # retain only the top_k most likely tokens, clamp others to have 0 probability
+
+    image_ori = input_img
+    image = image_transform(image_ori, resolution=config.dataset.params.resolution).to(device)
+    image = image.unsqueeze(0)
+    image_tokens = vq_model.get_code(image) + len(uni_prompting.text_tokenizer)
+
+    question = input_text
+    input_ids = uni_prompting.text_tokenizer(['USER: \n' + question + ' ASSISTANT:'])[
+        'input_ids']
+    input_ids = torch.tensor(input_ids).to(device)
+
+    input_ids = torch.cat([
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|mmu|>']).to(device),
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
+        image_tokens,
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|sot|>']).to(device),
+        input_ids
+    ], dim=1).long()
+
+    attention_mask = create_attention_mask_for_mmu(input_ids.to(device),
+                                                   eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']))
+
+    cont_toks_list = model.mmu_generate(input_ids, attention_mask=attention_mask,
+                                        max_new_tokens=100, top_k=top_k,
+                                        eot_token=uni_prompting.sptids_dict['<|eot|>'])
+
+    cont_toks_list = torch.stack(cont_toks_list).squeeze()[None]
+
+    output_text = uni_prompting.text_tokenizer.batch_decode(cont_toks_list, skip_special_tokens=True)
+
+    output_text = output_text[0].strip()
+
+    chat_history.append((input_text, output_text))
+
+    return "", chat_history
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""
+            <h1 class="display-2 fw-bold title">
+              <a style="color: #70a8dc;">S</a><a style="color: #6fb051;">h</a><a style="color: #e06766;">o</a><a style="color: #f7b26b;">w</a>-o
+            </h1>
+            <p>This is the official Gradio demo for the Show-o model, a unified model that can do multimodal understanding and generation.</p>
+
+            <strong>Paper:</strong> <a href="https://arxiv.org/abs/2408.12528" target="_blank">Show-o: One Single Transformer To Unify Multimodal Understanding and Generation </a>
+            <br/>
+            <strong>Project Website:</strong> <a href="https://showlab.github.io/Show-o/" target="_blank">Show-o Website</a>
+            <br/>
+            <strong>Code and Models:</strong> <a href="https://github.com/showlab/Show-o" target="_blank">GitHub</a>
+            <br/>
+            <br/>
+        """)
+
+    with gr.Row():
+        with gr.Column():
+            text_prompt_t2i = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image generation."
+            )
+            guidance_scale_t2i = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_t2i = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=18
+            )
+        generated_img_t2i = gr.Image(
+            label="Output image"
+        )
+    examples_t2i = gr.Examples(
+        label="Text to image generation examples",
+        examples=[
+            "A dynamic scene of a rally car race.",
+            "Paper artwork, layered paper, colorful Chinese dragon surrounded by clouds.",
+            "Pixel art character riding a dragon through the clouds.",
+        ],
+        inputs=text_prompt_t2i,
+    )
+    submit_btn_t2i = gr.Button("Generate: Text-to-image")
+    submit_btn_t2i.click(text_to_image_generation,
+                         [text_prompt_t2i, guidance_scale_t2i, generation_timesteps_t2i],
+                         [generated_img_t2i])
+
+    with gr.Row():
+        inpainting_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+        )
+        inpainting_input_mask = gr.Image(
+            label="Inpainting mask",
+            image_mode="L",
+            type="pil",
+        )
+
+        with gr.Column():
+            text_prompt_inpainting = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image inpainting."
+            )
+            guidance_scale_inpainting = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_inpainting = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+        generated_img_inpainting = gr.Image(
+            label="Output image"
+        )
+    examples_inpainting = gr.Examples(
+        label="Text-guided inpainting examples",
+        examples=[
+            [
+                "a blue sports car with sleek curves and tinted windows, parked on a bustling city street.",
+                Image.open("./inpainting_validation/bus.jpg").convert("RGB"),
+                Image.open("./inpainting_validation/bus_mask.webp").convert("L"),
+            ],
+            [
+                "a clear, shallow river with some vibrant flowers in it.",
+                Image.open("./inpainting_validation/train.jpg").convert("RGB"),
+                Image.open("./inpainting_validation/train_mask.webp").convert("L"),
+            ],
+        ],
+        inputs=[text_prompt_inpainting, inpainting_input_img, inpainting_input_mask],
+    )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Inpainting")
+    submit_btn_inpainting.click(text_guided_inpainting,
+                                [text_prompt_inpainting, inpainting_input_img, inpainting_input_mask,
+                                 guidance_scale_inpainting, generation_timesteps_inpainting],
+                                [generated_img_inpainting])
+
+    with gr.Row():
+        extra_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+            image_mode="RGB",
+        )
+
+        with gr.Column():
+            text_prompt_extrapolation = gr.Textbox(
+                label="Text prompt",
+                lines=1,
+                placeholder="Input the text prompt here for image extrapolation."
+            )
+            guidance_scale_extrapolation = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_extrapolation = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+            left_extrapolation = gr.Slider(
+                label="left extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+            right_extrapolation = gr.Slider(
+                label="right extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+        generated_img_extrapolation = gr.Image(
+            label="Output image"
+        )
+    examples_extra = gr.Examples(
+        label="Text-guided extrapolation examples",
+        examples=[
+            [
+                Image.open("./inpainting_validation/wukong2.jpg").convert("RGB"),
+                "the continuous mountain ranges and jungles, with meandering rivers occasionally appearing.",
+                2,
+                2,
+            ],
+            [
+                Image.open("./inpainting_validation/alpine_lake.jpg").convert("RGB"),
+                "a serene natural landscape featuring a clear, blue lake surrounded by lush green trees.",
+                2,
+                2,
+            ],
+        ],
+        inputs=[extra_input_img, text_prompt_extrapolation, left_extrapolation, right_extrapolation],
+    )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Extrapolation")
+    submit_btn_inpainting.click(text_guided_extrapolation,
+                                [extra_input_img, text_prompt_extrapolation, left_extrapolation, right_extrapolation,
+                                 guidance_scale_extrapolation, generation_timesteps_extrapolation],
+                                [generated_img_extrapolation])
+    with gr.Row():
+        with gr.Row():
+            chat_input_img = gr.Image(
+                label="Input image",
+                type="pil",
+                image_mode="RGB",
+            )
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            msg = gr.Textbox(label="Press Enter to send a message for chat")
+            clear = gr.ClearButton([msg, chatbot])
+    msg.submit(multimodal_understanding, [chat_input_img, msg, chatbot], [msg, chatbot])
+
+demo.launch()

configs/showo_demo.yaml

@@ -0,0 +1,49 @@
+wandb:
+  entity: null
+#  run_id: askkz9i2
+  resume: 'auto'
+
+experiment:
+    project: "demo"
+    name: "show-o-demo"
+    output_dir: "show-o-demo"
+
+model:
+    vq_model:
+        type: "magvitv2"
+        vq_model_name: "showlab/magvitv2"
+
+    showo:
+        pretrained_model_path: "showlab/show-o"
+        w_clip_vit: False
+        vocab_size: 58498
+        llm_vocab_size: 50295
+        llm_model_path: 'microsoft/phi-1_5'
+        codebook_size: 8192
+        num_vq_tokens: 256
+
+    gradient_checkpointing: True
+    enable_xformers_memory_efficient_attention: True
+
+
+dataset:
+    gen_type: "t2i"
+    und_type: "large_cap"
+    params:
+        batch_size: ${training.batch_size}
+        shuffle_buffer_size: 1000
+        num_workers: 32
+        resolution: 256
+        pin_memory: True
+        persistent_workers: True
+
+    preprocessing:
+        max_seq_length: 128
+        resolution: 256
+        center_crop: False
+        random_flip: False
+
+training:
+    gradient_accumulation_steps: 1
+    cond_dropout_prob: 0.1
+    batch_size: 20

configs/showo_demo_w_clip_vit.yaml

@@ -0,0 +1,49 @@
+wandb:
+  entity: null
+#  run_id: askkz9i2
+  resume: 'auto'
+
+experiment:
+    project: "demo"
+    name: "show-o-demo"
+    output_dir: "show-o-demo"
+
+model:
+    vq_model:
+        type: "magvitv2"
+        vq_model_name: "showlab/magvitv2"
+
+    showo:
+        pretrained_model_path: "showlab/show-o-w-clip-vit"
+        w_clip_vit: True
+        vocab_size: 58498
+        llm_vocab_size: 50295
+        llm_model_path: 'microsoft/phi-1_5'
+        codebook_size: 8192
+        num_vq_tokens: 256
+
+    gradient_checkpointing: True
+    enable_xformers_memory_efficient_attention: True
+
+
+dataset:
+    gen_type: "t2i"
+    und_type: "large_cap"
+    params:
+        batch_size: ${training.batch_size}
+        shuffle_buffer_size: 1000
+        num_workers: 32
+        resolution: 256
+        pin_memory: True
+        persistent_workers: True
+
+    preprocessing:
+        max_seq_length: 128
+        resolution: 256
+        center_crop: False
+        random_flip: False
+
+training:
+    gradient_accumulation_steps: 1
+    cond_dropout_prob: 0.1
+    batch_size: 20

gradio/app.py

@@ -0,0 +1,488 @@
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "False"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+import numpy as np
+import gradio as gr
+import torch
+from PIL import Image
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+import torch.nn.functional as F
+from transformers import CLIPImageProcessor
+
+import sys
+sys.path.insert(0, ".")
+from training import conversation as conversation_lib
+from prompting_utils import UniversalPrompting, create_attention_mask_predict_next, create_attention_mask_for_mmu
+from training_utils import image_transform
+from models import Showo, MAGVITv2, get_mask_chedule, CLIPVisionTower
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+conversation_lib.default_conversation = conversation_lib.conv_templates["phi1.5"]
+SYSTEM_PROMPT = "A chat between a curious user and an artificial intelligence assistant. " \
+                "The assistant gives helpful, detailed, and polite answers to the user's questions."
+SYSTEM_PROMPT_LEN = 28
+
+
+config = OmegaConf.load("configs/showo_demo.yaml")
+tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                   special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>",
+                                                   "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                   ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+vq_model = MAGVITv2()
+vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+vq_model.requires_grad_(False)
+vq_model.eval()
+
+model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+model.eval()
+mask_token_id = model.config.mask_token_id
+
+
+def text_to_image_generation(input_text, guidance_scale, generation_timesteps):
+    prompts = [input_text]
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    image_tokens = torch.ones((len(prompts), config.model.showo.num_vq_tokens),
+                              dtype=torch.long, device=device) * mask_token_id
+
+    input_ids, _ = uni_prompting((prompts, image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompts), image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def text_guided_inpainting(input_text, inpainting_image, inpainting_mask, guidance_scale, generation_timesteps):
+    prompt = [input_text]
+
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    inpainting_image = image_transform(inpainting_image, resolution=config.dataset.params.resolution).to(device)
+    inpainting_mask = image_transform(inpainting_mask, resolution=config.dataset.params.resolution, normalize=False)
+
+    inpainting_image = inpainting_image.unsqueeze(0).repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask = inpainting_mask.unsqueeze(0).to(device)
+    inpainting_mask = F.interpolate(inpainting_mask, size=config.dataset.params.resolution // 16, mode='bicubic')
+    inpainting_mask = inpainting_mask.repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask[inpainting_mask < 0.5] = 0
+    inpainting_mask[inpainting_mask >= 0.5] = 1
+
+    inpainting_mask = inpainting_mask.reshape(config.training.batch_size, -1)
+    inpainting_mask = inpainting_mask.to(torch.bool)
+
+    inpainting_image_tokens = vq_model.get_code(inpainting_image) + len(uni_prompting.text_tokenizer)
+    inpainting_image_tokens[inpainting_mask] = mask_token_id
+
+    input_ids, _ = uni_prompting((prompt, inpainting_image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompt), inpainting_image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def text_guided_extrapolation(input_img, input_text, left_ext, right_ext, guidance_scale, generation_timesteps):
+    config.offset = 0
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    extra_direction = ['right'] * int(right_ext) + ['left'] * int(left_ext)
+    prompt = [input_text] * len(extra_direction)
+    W = config.dataset.params.resolution // 16
+    for id, (prt, direction) in enumerate(zip(prompt, extra_direction)):
+        prt = [prt] * config.training.batch_size
+        if id == 0:
+            # extrapolation_image = Image.open(config.image_path).convert("RGB")
+            extrapolation_image = input_img
+            extrapolation_image = image_transform(extrapolation_image,
+                                                  resolution=config.dataset.params.resolution).to(device)
+
+            B, _, _ = extrapolation_image.shape
+            extrapolation_image = extrapolation_image.unsqueeze(0)
+            extrapolation_image_tokens = vq_model.get_code(extrapolation_image) + len(uni_prompting.text_tokenizer)
+            extrapolation_image_tokens = extrapolation_image_tokens.reshape(1,
+                                                                            config.dataset.params.resolution // 16,
+                                                                            config.dataset.params.resolution // 16)
+            extrapolation_image_tokens = extrapolation_image_tokens.repeat(config.training.batch_size, 1, 1)
+        else:
+
+            extrapolation_image_tokens = gen_token_ids + len(uni_prompting.text_tokenizer)
+
+        image_left_part = extrapolation_image_tokens[:, :, :-(W // 2 - config.offset)] - len(
+            uni_prompting.text_tokenizer)
+        image_right_part = extrapolation_image_tokens[:, :, W // 2 - config.offset:] - len(uni_prompting.text_tokenizer)
+        image_up_part = extrapolation_image_tokens[:, :-(W // 2 - config.offset), :] - len(uni_prompting.text_tokenizer)
+        image_down_part = extrapolation_image_tokens[:, W // 2 - config.offset:, :] - len(uni_prompting.text_tokenizer)
+
+        if direction in ['left', 'right']:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset),
+                                             dtype=torch.int64, device=device) + mask_token_id
+        else:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset,
+                                              config.dataset.params.resolution // 16),
+                                             dtype=torch.int64, device=device) + mask_token_id
+
+        if direction == 'left':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :, :W // 2 - config.offset]], dim=-1)
+        elif direction == 'right':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, :, -(W // 2 - config.offset):], extrapolation_mask], dim=-1)
+        elif direction == 'up':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :W // 2 - config.offset, :]], dim=-2)
+        else:
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, -(W // 2 - config.offset):, :], extrapolation_mask], dim=-2)
+
+        extrapolation_image_tokens = extrapolation_image_tokens.reshape(config.training.batch_size, -1)
+
+        input_ids, _ = uni_prompting((prt, extrapolation_image_tokens), 't2i_gen')
+
+        if config.training.guidance_scale > 0:
+            uncond_input_ids, _ = uni_prompting(([''] * len(prt), extrapolation_image_tokens), 't2i_gen')
+            attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+        else:
+            attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+            uncond_input_ids = None
+
+        if config.get("mask_schedule", None) is not None:
+            schedule = config.mask_schedule.schedule
+            args = config.mask_schedule.get("params", {})
+            mask_schedule = get_mask_chedule(schedule, **args)
+        else:
+            mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+        with torch.no_grad():
+            gen_token_ids = model.t2i_generate(
+                input_ids=input_ids,
+                uncond_input_ids=uncond_input_ids,
+                attention_mask=attention_mask,
+                guidance_scale=config.training.guidance_scale,
+                temperature=config.training.get("generation_temperature", 1.0),
+                timesteps=config.training.generation_timesteps,
+                noise_schedule=mask_schedule,
+                noise_type=config.training.get("noise_type", "mask"),
+                seq_len=config.model.showo.num_vq_tokens,
+                uni_prompting=uni_prompting,
+                config=config,
+            )
+
+        gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+        gen_token_ids = gen_token_ids.reshape(config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16)
+        if direction == 'left':
+            gen_token_ids = torch.cat([gen_token_ids, image_right_part], dim=-1)
+        elif direction == 'right':
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-1)
+        elif direction == 'up':
+            gen_token_ids = torch.cat([gen_token_ids, image_down_part], dim=-2)
+        else:
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-2)
+
+    _, h, w = gen_token_ids.shape
+    gen_token_ids = gen_token_ids.reshape(config.training.batch_size, -1)
+    images = vq_model.decode_code(gen_token_ids, shape=(h, w))
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def multimodal_understanding(input_img, input_text, chat_history):
+    top_k = 1  # retain only the top_k most likely tokens, clamp others to have 0 probability
+
+    image_ori = input_img
+    image = image_transform(image_ori, resolution=config.dataset.params.resolution).to(device)
+    image = image.unsqueeze(0)
+    image_tokens = vq_model.get_code(image) + len(uni_prompting.text_tokenizer)
+
+    question = input_text
+    input_ids = uni_prompting.text_tokenizer(['USER: \n' + question + ' ASSISTANT:'])[
+        'input_ids']
+    input_ids = torch.tensor(input_ids).to(device)
+
+    input_ids = torch.cat([
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|mmu|>']).to(device),
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
+        image_tokens,
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
+        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|sot|>']).to(device),
+        input_ids
+    ], dim=1).long()
+
+    attention_mask = create_attention_mask_for_mmu(input_ids.to(device),
+                                                   eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']))
+
+    cont_toks_list = model.mmu_generate(input_ids, attention_mask=attention_mask,
+                                        max_new_tokens=100, top_k=top_k,
+                                        eot_token=uni_prompting.sptids_dict['<|eot|>'])
+
+    cont_toks_list = torch.stack(cont_toks_list).squeeze()[None]
+
+    output_text = uni_prompting.text_tokenizer.batch_decode(cont_toks_list, skip_special_tokens=True)
+
+    output_text = output_text[0].strip()
+
+    chat_history.append((input_text, output_text))
+
+    return "", chat_history
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""
+            <h1 class="display-2 fw-bold title">
+              <a style="color: #70a8dc;">S</a><a style="color: #6fb051;">h</a><a style="color: #e06766;">o</a><a style="color: #f7b26b;">w</a>-o
+            </h1>
+            <p>This is the official Gradio demo for the Show-o model, a unified model that can do multimodal understanding and generation.</p>
+
+            <strong>Paper:</strong> <a href="https://arxiv.org/abs/2408.12528" target="_blank">Show-o: One Single Transformer To Unify Multimodal Understanding and Generation </a>
+            <br/>
+            <strong>Project Website:</strong> <a href="https://showlab.github.io/Show-o/" target="_blank">Show-o Website</a>
+            <br/>
+            <strong>Code and Models:</strong> <a href="https://github.com/showlab/Show-o" target="_blank">GitHub</a>
+            <br/>
+            <br/>
+        """)
+
+    with gr.Row():
+        with gr.Column():
+            text_prompt_t2i = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image generation."
+            )
+            guidance_scale_t2i = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_t2i = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=18
+            )
+        generated_img_t2i = gr.Image(
+            label="Output image"
+        )
+    submit_btn_t2i = gr.Button("Generate: Text-to-image")
+    submit_btn_t2i.click(text_to_image_generation,
+                         [text_prompt_t2i, guidance_scale_t2i, generation_timesteps_t2i],
+                         [generated_img_t2i])
+
+    with gr.Row():
+        inpainting_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+        )
+        inpainting_input_mask = gr.Image(
+            label="Inpainting mask",
+            image_mode="L",
+            type="pil",
+        )
+
+        with gr.Column():
+            text_prompt_inpainting = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image inpainting."
+            )
+            guidance_scale_inpainting = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_inpainting = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+        generated_img_inpainting = gr.Image(
+            label="Output image"
+        )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Inpainting")
+    submit_btn_inpainting.click(text_guided_inpainting,
+                                [text_prompt_inpainting, inpainting_input_img, inpainting_input_mask,
+                                 guidance_scale_inpainting, generation_timesteps_inpainting],
+                                [generated_img_inpainting])
+
+    with gr.Row():
+        extra_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+            image_mode="RGB",
+        )
+
+        with gr.Column():
+            text_prompt_extrapolation = gr.Textbox(
+                label="Text prompt",
+                lines=1,
+                placeholder="Input the text prompt here for image extrapolation."
+            )
+            guidance_scale_extrapolation = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_extrapolation = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+            left_extrapolation = gr.Slider(
+                label="left extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+            right_extrapolation = gr.Slider(
+                label="right extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+        generated_img_extrapolation = gr.Image(
+            label="Output image"
+        )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Extrapolation")
+    submit_btn_inpainting.click(text_guided_extrapolation,
+                                [extra_input_img, text_prompt_extrapolation, left_extrapolation, right_extrapolation,
+                                 guidance_scale_extrapolation, generation_timesteps_extrapolation],
+                                [generated_img_extrapolation])
+
+    with gr.Row():
+        with gr.Row():
+            chat_input_img = gr.Image(
+                label="Input image",
+                type="pil",
+                image_mode="RGB",
+            )
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            msg = gr.Textbox(label="Press Enter to send a message for chat")
+            clear = gr.ClearButton([msg, chatbot])
+    msg.submit(multimodal_understanding, [chat_input_img, msg, chatbot], [msg, chatbot])
+
+demo.launch()

gradio/app_gradio.py

@@ -0,0 +1,281 @@
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "False"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+import tempfile
+from share_btn import share_js, save_js
+import gradio as gr
+from PIL import Image
+import torch
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+
+from models import Showo, MAGVITv2, get_mask_chedule
+from prompting_utils import UniversalPrompting, create_attention_mask_predict_next
+
+
+# Prepare model
+config = OmegaConf.load("configs/showo_demo.yaml")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                   special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                   ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+vq_model = MAGVITv2(config.model.vq_model.type)
+vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+vq_model.requires_grad_(False)
+vq_model.eval()
+
+model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+model.eval()
+
+mask_token_id = model.config.mask_token_id
+
+
+css = """
+    #chatbot { min-height: 300px; }
+    #save-btn {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    #save-btn:hover {
+        background-image: linear-gradient(to right bottom, rgba(110,197,224, 0.9), rgba(138,211,194, 1.0));
+    }
+    #share-btn {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    #share-btn:hover {
+        background-image: linear-gradient(to right bottom, rgba(110,197,224, 0.9), rgba(138,211,194, 1.0));
+    }
+    #gallery { z-index: 999999; }
+    #gallery img:hover {transform: scale(2.3); z-index: 999999; position: relative; padding-right: 30%; padding-bottom: 30%;}
+    #gallery button img:hover {transform: none; z-index: 999999; position: relative; padding-right: 0; padding-bottom: 0;}
+    @media (hover: none) {
+        #gallery img:hover {transform: none; z-index: 999999; position: relative; padding-right: 0; 0;}
+    }
+    .html2canvas-container { width: 3000px !important; height: 3000px !important; }
+"""
+
+
+def upload_image(state, image_input):
+    conversation = state[0]
+    chat_history = state[1]
+    input_image = Image.open(image_input.name).resize(
+        (224, 224)).convert('RGB')
+    input_image.save(image_input.name)  # Overwrite with smaller image.
+    conversation += [(f'<img src="./file={image_input.name}" style="display: inline-block;">', "")]
+    return [conversation, chat_history + [input_image, ""]], conversation
+
+
+def reset():
+    return [[], []], []
+
+
+def reset_last(state):
+    conversation = state[0][:-1]
+    chat_history = state[1][:-2]
+    return [conversation, chat_history], conversation
+
+
+def save_image_to_local(image: Image.Image):
+    filename = next(tempfile._get_candidate_names()) + '.png'
+    image.save(filename)
+    return filename
+
+
+def text_to_image_generation(input_text, state, guidance_scale, generation_timesteps):
+    prompts = [input_text]
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    image_tokens = torch.ones((len(prompts), config.model.showo.num_vq_tokens),
+                              dtype=torch.long, device=device) * mask_token_id
+
+    input_ids, _ = uni_prompting((prompts, image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompts), image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+    pil_images = [Image.fromarray(image) for image in images]
+
+    wandb_images = [wandb.Image(image, caption=prompts[i]) for i, image in enumerate(pil_images)]
+    wandb.log({"generated_images": wandb_images}, step=step)
+
+
+def generate_for_prompt(input_text, state, ret_scale_factor, num_words, temperature):
+    g_cuda = torch.Generator(device='cuda').manual_seed(1337)
+
+    # Ignore empty inputs.
+    if len(input_text) == 0:
+        return state, state[0], gr.update(visible=True)
+
+    input_prompt = 'Q: ' + input_text + '\nA:'
+    conversation = state[0]
+    chat_history = state[1]
+    print('Generating for', chat_history, flush=True)
+
+    # If an image was uploaded, prepend it to the model.
+    model_inputs = chat_history
+    model_inputs.append(input_prompt)
+    # Remove empty text.
+    model_inputs = [s for s in model_inputs if s != '']
+
+    top_p = 1.0
+    if temperature != 0.0:
+        top_p = 0.95
+
+    print('Running model.generate_for_images_and_texts with', model_inputs, flush=True)
+    model_outputs = model.generate_for_images_and_texts(model_inputs,
+                                                        num_words=max(num_words, 1), ret_scale_factor=ret_scale_factor, top_p=top_p,
+                                                        temperature=temperature, max_num_rets=1,
+                                                        num_inference_steps=50, generator=g_cuda)
+    print('model_outputs', model_outputs, ret_scale_factor, flush=True)
+
+    response = ''
+    text_outputs = []
+    for output_i, p in enumerate(model_outputs):
+        if type(p) == str:
+            if output_i > 0:
+                response += '<br/>'
+            # Remove the image tokens for output.
+            text_outputs.append(p.replace('[IMG0] [IMG1] [IMG2] [IMG3] [IMG4] [IMG5] [IMG6] [IMG7]', ''))
+            response += p
+            if len(model_outputs) > 1:
+                response += '<br/>'
+        elif type(p) == dict:
+            # Decide whether to generate or retrieve.
+            if p['decision'] is not None and p['decision'][0] == 'gen':
+                image = p['gen'][0][0]#.resize((224, 224))
+                filename = save_image_to_local(image)
+                response += f'<img src="./file={filename}" style="display: inline-block;"><p style="font-size: 12px; color: #555; margin-top: 0;">(Generated)</p>'
+            else:
+                image = p['ret'][0][0]#.resize((224, 224))
+                filename = save_image_to_local(image)
+                response += f'<img src="./file={filename}" style="display: inline-block;"><p style="font-size: 12px; color: #555; margin-top: 0;">(Retrieved)</p>'
+
+    chat_history = model_inputs + \
+        [' '.join([s for s in model_outputs if type(s) == str]) + '\n']
+    # Remove [RET] from outputs.
+    conversation.append((input_text, response.replace('[IMG0] [IMG1] [IMG2] [IMG3] [IMG4] [IMG5] [IMG6] [IMG7]', '')))
+
+    # Set input image to None.
+    print('state', state, flush=True)
+    print('updated state', [conversation, chat_history], flush=True)
+    return [conversation, chat_history], conversation, gr.update(visible=True), gr.update(visible=True)
+
+
+with gr.Blocks(css=css) as demo:
+    gr.HTML("""
+        <h1>🐟 GILL</h1>
+        <p>This is the official Gradio demo for the GILL model, a model that can process arbitrarily interleaved image and text inputs, and produce image and text outputs.</p>
+
+        <strong>Paper:</strong> <a href="https://arxiv.org/abs/2305.17216" target="_blank">Generating Images with Multimodal Language Models</a>
+        <br/>
+        <strong>Project Website:</strong> <a href="https://jykoh.com/gill" target="_blank">GILL Website</a>
+        <br/>
+        <strong>Code and Models:</strong> <a href="https://github.com/kohjingyu/gill" target="_blank">GitHub</a>
+        <br/>
+        <br/>
+
+        <strong>Tips:</strong>
+        <ul>
+        <li>Start by inputting either image or text prompts (or both) and chat with GILL to get image-and-text replies.</li>
+        <li>Tweak the level of sensitivity to images and text using the parameters on the right.</li>
+        <li>Check out cool conversations in the examples or community tab for inspiration and share your own!</li>
+        <li>If the model outputs a blank image, it is because Stable Diffusion's safety filter detected inappropriate content. Please try again with a different prompt.</li>
+        <li>Outputs may differ slightly from the paper due to slight implementation differences. For reproducing paper results, please use our <a href="https://github.com/kohjingyu/gill" target="_blank">official code</a>.</li>
+        <li>For faster inference without waiting in queue, you may duplicate the space and use your own GPU: <a href="https://huggingface.co/spaces/jykoh/gill?duplicate=true"><img style="display: inline-block; margin-top: 0em; margin-bottom: 0em" src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></li>
+        </ul>
+    """)
+
+    gr_state = gr.State([[], []])  # conversation, chat_history
+
+    with gr.Row():
+        with gr.Column(scale=0.7, min_width=500):
+            with gr.Row():
+                chatbot = gr.Chatbot(elem_id="chatbot", label="🐟 GILL Chatbot")
+            with gr.Row():
+                image_btn = gr.UploadButton("🖼️ Upload Image", file_types=["image"])
+
+                text_input = gr.Textbox(label="Message", placeholder="Type a message")
+
+                with gr.Column():
+                    submit_btn = gr.Button("Submit", interactive=True, variant="primary")
+                    clear_last_btn = gr.Button("Undo")
+                    clear_btn = gr.Button("Reset All")
+                    with gr.Row(visible=False) as save_group:
+                        save_button = gr.Button("💾 Save Conversation as .png", elem_id="save-btn")
+
+                    with gr.Row(visible=False) as share_group:
+                        share_button = gr.Button("🤗 Share to Community (opens new window)", elem_id="share-btn")
+
+        with gr.Column(scale=0.3, min_width=400):
+            ret_scale_factor = gr.Slider(minimum=0.0, maximum=3.0, value=1.3, step=0.1, interactive=True,
+                                         label="Frequency multiplier for returning images (higher means more frequent)")
+            gr_max_len = gr.Slider(minimum=1, maximum=64, value=32,
+                                   step=1, interactive=True, label="Max # of words")
+            gr_temperature = gr.Slider(
+                minimum=0.0, maximum=1.0, value=0.0, step=0.1, interactive=True, label="Temperature (0 for deterministic, higher for more randomness)")
+
+            gallery = gr.Gallery(
+                value=[Image.open(e) for e in examples], label="Example Conversations", show_label=True, elem_id="gallery",
+            ).style(grid=[2], height="auto")
+
+    text_input.submit(generate_for_prompt, [text_input, gr_state, ret_scale_factor,
+                      gr_max_len, gr_temperature], [gr_state, chatbot, share_group, save_group])
+    text_input.submit(lambda: "", None, text_input)  # Reset chatbox.
+
+    submit_btn.click(generate_for_prompt, [text_input, gr_state, ret_scale_factor,
+                     gr_max_len, gr_temperature], [gr_state, chatbot, share_group, save_group])
+    submit_btn.click(lambda: "", None, text_input)  # Reset chatbox.
+
+    image_btn.upload(upload_image, [gr_state, image_btn], [gr_state, chatbot])
+    clear_last_btn.click(reset_last, [gr_state], [gr_state, chatbot])
+    clear_btn.click(reset, [], [gr_state, chatbot])
+    share_button.click(None, [], [], _js=share_js)
+    save_button.click(None, [], [], _js=save_js)
+
+
+demo.queue(concurrency_count=1, api_open=False, max_size=16)
+demo.launch(debug=True, server_name="0.0.0.0")

gradio/app_w_clip.py

@@ -0,0 +1,559 @@
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "False"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+import numpy as np
+import gradio as gr
+import torch
+from PIL import Image
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+import torch.nn.functional as F
+from transformers import CLIPImageProcessor
+
+import sys
+sys.path.insert(0, ".")
+from training import conversation as conversation_lib
+from prompting_utils import UniversalPrompting, create_attention_mask_predict_next, create_attention_mask_for_mmu_vit
+from training_utils import image_transform
+from models import Showo, MAGVITv2, get_mask_chedule, CLIPVisionTower
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+conversation_lib.default_conversation = conversation_lib.conv_templates["phi1.5"]
+SYSTEM_PROMPT = "A chat between a curious user and an artificial intelligence assistant. " \
+                "The assistant gives helpful, detailed, and polite answers to the user's questions."
+SYSTEM_PROMPT_LEN = 28
+
+
+def load_discrete_checkpoint():
+    config = OmegaConf.load("configs/showo_demo.yaml")
+    tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+    uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                       special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>",
+                                                       "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                       ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+    vq_model = MAGVITv2()
+    vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+    vq_model.requires_grad_(False)
+    vq_model.eval()
+
+    model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+    model.eval()
+    mask_token_id = model.config.mask_token_id
+
+    return config, uni_prompting, tokenizer, vq_model, model, mask_token_id
+
+
+config_gen, uni_prompting_gen, tokenizer_gen, vq_model_gen, model_gen, mask_token_id = load_discrete_checkpoint()
+
+
+def load_continuous_checkpoint():
+    config = OmegaConf.load("configs/showo_demo_w_clip_vit.yaml")
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+    uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                       special_tokens=(
+                                       "<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>",
+                                       "<|v2v|>", "<|lvg|>"),
+                                       ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+    vision_tower_name = "openai/clip-vit-large-patch14-336"
+    vision_tower = CLIPVisionTower(vision_tower_name).to(device)
+    clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
+
+    model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+    model.eval()
+
+    return config, uni_prompting, tokenizer, model, vision_tower, clip_image_processor
+
+
+config_mmu = uni_prompting_mmu = tokenizer_mmu = model_mmu = vision_tower = clip_image_processor = None
+
+
+def text_to_image_generation(input_text, guidance_scale, generation_timesteps):
+    config, uni_prompting, tokenizer, vq_model, model = config_gen, uni_prompting_gen, tokenizer_gen, vq_model_gen, model_gen
+
+    prompts = [input_text]
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    image_tokens = torch.ones((len(prompts), config.model.showo.num_vq_tokens),
+                              dtype=torch.long, device=device) * mask_token_id
+
+    input_ids, _ = uni_prompting((prompts, image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompts), image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def text_guided_inpainting(input_text, inpainting_image, inpainting_mask, guidance_scale, generation_timesteps):
+    config, uni_prompting, tokenizer, vq_model, model = config_gen, uni_prompting_gen, tokenizer_gen, vq_model_gen, model_gen
+
+    prompt = [input_text]
+
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    inpainting_image = image_transform(inpainting_image, resolution=config.dataset.params.resolution).to(device)
+    inpainting_mask = image_transform(inpainting_mask, resolution=config.dataset.params.resolution, normalize=False)
+
+    inpainting_image = inpainting_image.unsqueeze(0).repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask = inpainting_mask.unsqueeze(0).to(device)
+    inpainting_mask = F.interpolate(inpainting_mask, size=config.dataset.params.resolution // 16, mode='bicubic')
+    inpainting_mask = inpainting_mask.repeat(config.training.batch_size, 1, 1, 1)
+
+    inpainting_mask[inpainting_mask < 0.5] = 0
+    inpainting_mask[inpainting_mask >= 0.5] = 1
+
+    inpainting_mask = inpainting_mask.reshape(config.training.batch_size, -1)
+    inpainting_mask = inpainting_mask.to(torch.bool)
+
+    inpainting_image_tokens = vq_model.get_code(inpainting_image) + len(uni_prompting.text_tokenizer)
+    inpainting_image_tokens[inpainting_mask] = mask_token_id
+
+    input_ids, _ = uni_prompting((prompt, inpainting_image_tokens), 't2i_gen')
+
+    if config.training.guidance_scale > 0:
+        uncond_input_ids, _ = uni_prompting(([''] * len(prompt), inpainting_image_tokens), 't2i_gen')
+        attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+    else:
+        attention_mask = create_attention_mask_predict_next(input_ids,
+                                                            pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                            soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                            eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                            rm_pad_in_image=True)
+        uncond_input_ids = None
+
+    if config.get("mask_schedule", None) is not None:
+        schedule = config.mask_schedule.schedule
+        args = config.mask_schedule.get("params", {})
+        mask_schedule = get_mask_chedule(schedule, **args)
+    else:
+        mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+    with torch.no_grad():
+        gen_token_ids = model.t2i_generate(
+            input_ids=input_ids,
+            uncond_input_ids=uncond_input_ids,
+            attention_mask=attention_mask,
+            guidance_scale=config.training.guidance_scale,
+            temperature=config.training.get("generation_temperature", 1.0),
+            timesteps=config.training.generation_timesteps,
+            noise_schedule=mask_schedule,
+            noise_type=config.training.get("noise_type", "mask"),
+            seq_len=config.model.showo.num_vq_tokens,
+            uni_prompting=uni_prompting,
+            config=config,
+        )
+
+    gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+    images = vq_model.decode_code(gen_token_ids)
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def text_guided_extrapolation(input_img, input_text, left_ext, right_ext, guidance_scale, generation_timesteps):
+    config, uni_prompting, tokenizer, vq_model, model = config_gen, uni_prompting_gen, tokenizer_gen, vq_model_gen, model_gen
+
+    config.offset = 0
+    config.training.batch_size = config.batch_size = 1
+    config.training.guidance_scale = config.guidance_scale = guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps = generation_timesteps
+
+    extra_direction = ['right'] * int(right_ext) + ['left'] * int(left_ext)
+    prompt = [input_text] * len(extra_direction)
+    W = config.dataset.params.resolution // 16
+    for id, (prt, direction) in enumerate(zip(prompt, extra_direction)):
+        prt = [prt] * config.training.batch_size
+        if id == 0:
+            # extrapolation_image = Image.open(config.image_path).convert("RGB")
+            extrapolation_image = input_img
+            extrapolation_image = image_transform(extrapolation_image,
+                                                  resolution=config.dataset.params.resolution).to(device)
+
+            B, _, _ = extrapolation_image.shape
+            extrapolation_image = extrapolation_image.unsqueeze(0)
+            extrapolation_image_tokens = vq_model.get_code(extrapolation_image) + len(uni_prompting.text_tokenizer)
+            extrapolation_image_tokens = extrapolation_image_tokens.reshape(1,
+                                                                            config.dataset.params.resolution // 16,
+                                                                            config.dataset.params.resolution // 16)
+            extrapolation_image_tokens = extrapolation_image_tokens.repeat(config.training.batch_size, 1, 1)
+        else:
+
+            extrapolation_image_tokens = gen_token_ids + len(uni_prompting.text_tokenizer)
+
+        image_left_part = extrapolation_image_tokens[:, :, :-(W // 2 - config.offset)] - len(
+            uni_prompting.text_tokenizer)
+        image_right_part = extrapolation_image_tokens[:, :, W // 2 - config.offset:] - len(uni_prompting.text_tokenizer)
+        image_up_part = extrapolation_image_tokens[:, :-(W // 2 - config.offset), :] - len(uni_prompting.text_tokenizer)
+        image_down_part = extrapolation_image_tokens[:, W // 2 - config.offset:, :] - len(uni_prompting.text_tokenizer)
+
+        if direction in ['left', 'right']:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset),
+                                             dtype=torch.int64, device=device) + mask_token_id
+        else:
+            extrapolation_mask = torch.zeros((config.training.batch_size,
+                                              config.dataset.params.resolution // 16 // 2 + config.offset,
+                                              config.dataset.params.resolution // 16),
+                                             dtype=torch.int64, device=device) + mask_token_id
+
+        if direction == 'left':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :, :W // 2 - config.offset]], dim=-1)
+        elif direction == 'right':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, :, -(W // 2 - config.offset):], extrapolation_mask], dim=-1)
+        elif direction == 'up':
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_mask, extrapolation_image_tokens[:, :W // 2 - config.offset, :]], dim=-2)
+        else:
+            extrapolation_image_tokens = torch.cat(
+                [extrapolation_image_tokens[:, -(W // 2 - config.offset):, :], extrapolation_mask], dim=-2)
+
+        extrapolation_image_tokens = extrapolation_image_tokens.reshape(config.training.batch_size, -1)
+
+        input_ids, _ = uni_prompting((prt, extrapolation_image_tokens), 't2i_gen')
+
+        if config.training.guidance_scale > 0:
+            uncond_input_ids, _ = uni_prompting(([''] * len(prt), extrapolation_image_tokens), 't2i_gen')
+            attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+        else:
+            attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+            uncond_input_ids = None
+
+        if config.get("mask_schedule", None) is not None:
+            schedule = config.mask_schedule.schedule
+            args = config.mask_schedule.get("params", {})
+            mask_schedule = get_mask_chedule(schedule, **args)
+        else:
+            mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+        with torch.no_grad():
+            gen_token_ids = model.t2i_generate(
+                input_ids=input_ids,
+                uncond_input_ids=uncond_input_ids,
+                attention_mask=attention_mask,
+                guidance_scale=config.training.guidance_scale,
+                temperature=config.training.get("generation_temperature", 1.0),
+                timesteps=config.training.generation_timesteps,
+                noise_schedule=mask_schedule,
+                noise_type=config.training.get("noise_type", "mask"),
+                seq_len=config.model.showo.num_vq_tokens,
+                uni_prompting=uni_prompting,
+                config=config,
+            )
+
+        gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+        gen_token_ids = gen_token_ids.reshape(config.training.batch_size,
+                                              config.dataset.params.resolution // 16,
+                                              config.dataset.params.resolution // 16)
+        if direction == 'left':
+            gen_token_ids = torch.cat([gen_token_ids, image_right_part], dim=-1)
+        elif direction == 'right':
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-1)
+        elif direction == 'up':
+            gen_token_ids = torch.cat([gen_token_ids, image_down_part], dim=-2)
+        else:
+            gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-2)
+
+    _, h, w = gen_token_ids.shape
+    gen_token_ids = gen_token_ids.reshape(config.training.batch_size, -1)
+    images = vq_model.decode_code(gen_token_ids, shape=(h, w))
+
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+
+    return images[0]
+
+
+def multimodal_understanding(input_img, input_text, chat_history):
+    global config_mmu, uni_prompting_mmu, tokenizer_mmu, model_mmu, vision_tower, clip_image_processor
+    if model_mmu is None:
+        config_mmu, uni_prompting_mmu, tokenizer_mmu, model_mmu, vision_tower, clip_image_processor = load_continuous_checkpoint()
+    config, uni_prompting, tokenizer, model = config_mmu, uni_prompting_mmu, tokenizer_mmu, model_mmu
+
+    image_ori = input_img
+    pixel_values = clip_image_processor.preprocess(image_ori, return_tensors="pt")["pixel_values"][0]
+    batch_size = 1
+    question = input_text
+    top_k = 1  # retain only the top_k most likely tokens, clamp others to have 0 probability
+
+    conv = conversation_lib.default_conversation.copy()
+    conv.append_message(conv.roles[0], question)
+    conv.append_message(conv.roles[1], None)
+    prompt_question = conv.get_prompt()
+    question_input = []
+    question_input.append(prompt_question.strip())
+
+    input_ids_system = [uni_prompting.text_tokenizer(SYSTEM_PROMPT, return_tensors="pt", padding="longest").input_ids
+                            for _ in range(batch_size)]
+    input_ids_system = torch.stack(input_ids_system, dim=0)
+    assert input_ids_system.shape[-1] == 28
+    input_ids_system = input_ids_system.to(device)
+    input_ids_system = input_ids_system[0]
+
+    input_ids = [uni_prompting.text_tokenizer(prompt, return_tensors="pt", padding="longest").input_ids
+                    for prompt in question_input]
+
+    input_ids = torch.stack(input_ids)
+    input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=uni_prompting.text_tokenizer.pad_token_id
+    )
+    input_ids = torch.tensor(input_ids).to(device).squeeze(0)
+    input_ids_llava = torch.cat([
+            (torch.ones(input_ids.shape[0], 1) *uni_prompting.sptids_dict['<|mmu|>']).to(device),
+            input_ids_system,
+            (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
+            # place your img embedding here
+            (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
+            input_ids,
+    ], dim=1).long()
+
+    images_embeddings = vision_tower(pixel_values[None])
+    images_embeddings = model.mm_projector(images_embeddings)
+
+    text_embeddings = model.showo.model.embed_tokens(input_ids_llava)
+
+    # Full input seq
+    part1 = text_embeddings[:, :2 + SYSTEM_PROMPT_LEN, :]
+    part2 = text_embeddings[:, 2 + SYSTEM_PROMPT_LEN:, :]
+    input_embeddings = torch.cat((part1, images_embeddings, part2), dim=1)
+
+    attention_mask_llava = create_attention_mask_for_mmu_vit(input_embeddings,
+                                                            system_prompt_len=SYSTEM_PROMPT_LEN)
+
+    cont_toks_list = model.mmu_generate(input_embeddings=input_embeddings,
+                                        attention_mask=attention_mask_llava[0].unsqueeze(0),
+                                        max_new_tokens=100,
+                                        top_k=top_k,
+                                        # eot_token=uni_prompting.sptids_dict['<|eot|>']
+                                        eot_token=tokenizer.eos_token_id
+                                        )
+
+    cont_toks_list = torch.stack(cont_toks_list).squeeze()[None]
+
+    output_text = uni_prompting.text_tokenizer.batch_decode(cont_toks_list, skip_special_tokens=True)
+
+    output_text = output_text[0].strip()
+
+    chat_history.append((input_text, output_text))
+
+    return "", chat_history
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""
+            <h1 class="display-2 fw-bold title">
+              <a style="color: #70a8dc;">S</a><a style="color: #6fb051;">h</a><a style="color: #e06766;">o</a><a style="color: #f7b26b;">w</a>-o
+            </h1>
+            <p>This is the official Gradio demo for the Show-o model, a unified model that can do multimodal understanding and generation.</p>
+
+            <strong>Paper:</strong> <a href="https://arxiv.org/abs/2408.12528" target="_blank">Show-o: One Single Transformer To Unify Multimodal Understanding and Generation </a>
+            <br/>
+            <strong>Project Website:</strong> <a href="https://showlab.github.io/Show-o/" target="_blank">Show-o Website</a>
+            <br/>
+            <strong>Code and Models:</strong> <a href="https://github.com/showlab/Show-o" target="_blank">GitHub</a>
+            <br/>
+            <br/>
+        """)
+
+    with gr.Row():
+        with gr.Column():
+            text_prompt_t2i = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image generation."
+            )
+            guidance_scale_t2i = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_t2i = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=18
+            )
+        generated_img_t2i = gr.Image(
+            label="Output image"
+        )
+    submit_btn_t2i = gr.Button("Generate: Text-to-image")
+    submit_btn_t2i.click(text_to_image_generation,
+                         [text_prompt_t2i, guidance_scale_t2i, generation_timesteps_t2i],
+                         [generated_img_t2i])
+
+    with gr.Row():
+        inpainting_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+        )
+        inpainting_input_mask = gr.Image(
+            label="Inpainting mask",
+            image_mode="L",
+            type="pil",
+        )
+
+        with gr.Column():
+            text_prompt_inpainting = gr.Textbox(
+                label="Text prompt",
+                lines=2,
+                placeholder="Input the text prompt here for image inpainting."
+            )
+            guidance_scale_inpainting = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_inpainting = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+        generated_img_inpainting = gr.Image(
+            label="Output image"
+        )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Inpainting")
+    submit_btn_inpainting.click(text_guided_inpainting,
+                                [text_prompt_inpainting, inpainting_input_img, inpainting_input_mask,
+                                 guidance_scale_inpainting, generation_timesteps_inpainting],
+                                [generated_img_inpainting])
+
+    with gr.Row():
+        extra_input_img = gr.Image(
+            label="Input image",
+            type="pil",
+            image_mode="RGB",
+        )
+
+        with gr.Column():
+            text_prompt_extrapolation = gr.Textbox(
+                label="Text prompt",
+                lines=1,
+                placeholder="Input the text prompt here for image extrapolation."
+            )
+            guidance_scale_extrapolation = gr.Slider(
+                label="guidance scale",
+                minimum=0,
+                maximum=5,
+                step=0.05,
+                value=1.75
+            )
+            generation_timesteps_extrapolation = gr.Slider(
+                label="timesteps",
+                minimum=1,
+                maximum=30,
+                step=1,
+                value=16
+            )
+            left_extrapolation = gr.Slider(
+                label="left extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+            right_extrapolation = gr.Slider(
+                label="right extrapolation",
+                minimum=0,
+                maximum=5,
+                step=1,
+                value=1
+            )
+        generated_img_extrapolation = gr.Image(
+            label="Output image"
+        )
+    submit_btn_inpainting = gr.Button("Generate: Text-guided Extrapolation")
+    submit_btn_inpainting.click(text_guided_extrapolation,
+                                [extra_input_img, text_prompt_extrapolation, left_extrapolation, right_extrapolation,
+                                 guidance_scale_extrapolation, generation_timesteps_extrapolation],
+                                [generated_img_extrapolation])
+
+    with gr.Row():
+        with gr.Row():
+            chat_input_img = gr.Image(
+                label="Input image",
+                type="pil",
+                image_mode="RGB",
+            )
+        with gr.Column():
+            chatbot = gr.Chatbot()
+            msg = gr.Textbox(label="Press Enter to send a message for chat")
+            clear = gr.ClearButton([msg, chatbot])
+    msg.submit(multimodal_understanding, [chat_input_img, msg, chatbot], [msg, chatbot])
+
+demo.launch()

gradio/share_btn.py

@@ -0,0 +1,113 @@
+# Modified from https://huggingface.co/spaces/haoheliu/audioldm-text-to-audio-generation/blob/79681cd8cb235160a27cdd100673346eb1784e53/share_btn.py
+
+community_icon_html = """<svg id="share-btn-share-icon" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32">
+    <path d="M20.6081 3C21.7684 3 22.8053 3.49196 23.5284 4.38415C23.9756 4.93678 24.4428 5.82749 24.4808 7.16133C24.9674 7.01707 25.4353 6.93643 25.8725 6.93643C26.9833 6.93643 27.9865 7.37587 28.696 8.17411C29.6075 9.19872 30.0124 10.4579 29.8361 11.7177C29.7523 12.3177 29.5581 12.8555 29.2678 13.3534C29.8798 13.8646 30.3306 14.5763 30.5485 15.4322C30.719 16.1032 30.8939 17.5006 29.9808 18.9403C30.0389 19.0342 30.0934 19.1319 30.1442 19.2318C30.6932 20.3074 30.7283 21.5229 30.2439 22.6548C29.5093 24.3704 27.6841 25.7219 24.1397 27.1727C21.9347 28.0753 19.9174 28.6523 19.8994 28.6575C16.9842 29.4379 14.3477 29.8345 12.0653 29.8345C7.87017 29.8345 4.8668 28.508 3.13831 25.8921C0.356375 21.6797 0.754104 17.8269 4.35369 14.1131C6.34591 12.058 7.67023 9.02782 7.94613 8.36275C8.50224 6.39343 9.97271 4.20438 12.4172 4.20438H12.4179C12.6236 4.20438 12.8314 4.2214 13.0364 4.25468C14.107 4.42854 15.0428 5.06476 15.7115 6.02205C16.4331 5.09583 17.134 4.359 17.7682 3.94323C18.7242 3.31737 19.6794 3 20.6081 3ZM20.6081 5.95917C20.2427 5.95917 19.7963 6.1197 19.3039 6.44225C17.7754 7.44319 14.8258 12.6772 13.7458 14.7131C13.3839 15.3952 12.7655 15.6837 12.2086 15.6837C11.1036 15.6837 10.2408 14.5497 12.1076 13.1085C14.9146 10.9402 13.9299 7.39584 12.5898 7.1776C12.5311 7.16799 12.4731 7.16355 12.4172 7.16355C11.1989 7.16355 10.6615 9.33114 10.6615 9.33114C10.6615 9.33114 9.0863 13.4148 6.38031 16.206C3.67434 18.998 3.5346 21.2388 5.50675 24.2246C6.85185 26.2606 9.42666 26.8753 12.0653 26.8753C14.8021 26.8753 17.6077 26.2139 19.1799 25.793C19.2574 25.7723 28.8193 22.984 27.6081 20.6107C27.4046 20.212 27.0693 20.0522 26.6471 20.0522C24.9416 20.0522 21.8393 22.6726 20.5057 22.6726C20.2076 22.6726 19.9976 22.5416 19.9116 22.222C19.3433 20.1173 28.552 19.2325 27.7758 16.1839C27.639 15.6445 27.2677 15.4256 26.746 15.4263C24.4923 15.4263 19.4358 19.5181 18.3759 19.5181C18.2949 19.5181 18.2368 19.4937 18.2053 19.4419C17.6743 18.557 17.9653 17.9394 21.7082 15.6009C25.4511 13.2617 28.0783 11.8545 26.5841 10.1752C26.4121 9.98141 26.1684 9.8956 25.8725 9.8956C23.6001 9.89634 18.2311 14.9403 18.2311 14.9403C18.2311 14.9403 16.7821 16.496 15.9057 16.496C15.7043 16.496 15.533 16.4139 15.4169 16.2112C14.7956 15.1296 21.1879 10.1286 21.5484 8.06535C21.7928 6.66715 21.3771 5.95917 20.6081 5.95917Z" fill="#FF9D00"></path>
+    <path d="M5.50686 24.2246C3.53472 21.2387 3.67446 18.9979 6.38043 16.206C9.08641 13.4147 10.6615 9.33111 10.6615 9.33111C10.6615 9.33111 11.2499 6.95933 12.59 7.17757C13.93 7.39581 14.9139 10.9401 12.1069 13.1084C9.29997 15.276 12.6659 16.7489 13.7459 14.713C14.8258 12.6772 17.7747 7.44316 19.304 6.44221C20.8326 5.44128 21.9089 6.00204 21.5484 8.06532C21.188 10.1286 14.795 15.1295 15.4171 16.2118C16.0391 17.2934 18.2312 14.9402 18.2312 14.9402C18.2312 14.9402 25.0907 8.49588 26.5842 10.1752C28.0776 11.8545 25.4512 13.2616 21.7082 15.6008C17.9646 17.9393 17.6744 18.557 18.2054 19.4418C18.7372 20.3266 26.9998 13.1351 27.7759 16.1838C28.5513 19.2324 19.3434 20.1173 19.9117 22.2219C20.48 24.3274 26.3979 18.2382 27.6082 20.6107C28.8193 22.9839 19.2574 25.7722 19.18 25.7929C16.0914 26.62 8.24723 28.3726 5.50686 24.2246Z" fill="#FFD21E"></path>
+</svg>"""
+
+loading_icon_html = """<svg id="share-btn-loading-icon" style="display:none;" class="animate-spin"
+   style="color: #ffffff; 
+"
+   xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" fill="none" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 24 24"><circle style="opacity: 0.25;" cx="12" cy="12" r="10" stroke="white" stroke-width="4"></circle><path style="opacity: 0.75;" fill="white" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg>"""
+
+share_js = """
+async () => {
+    const html2canvas = (await import('https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.esm.js')).default;
+    async function uploadFile(file) {
+        console.log(file.type)
+        const UPLOAD_URL = 'https://huggingface.co/uploads';
+        const response = await fetch(UPLOAD_URL, {
+            method: 'POST',
+            headers: {
+                'Content-Type': file.type,
+                'X-Requested-With': 'XMLHttpRequest',
+            },
+            body: file, /// <- File inherits from Blob
+        });
+        const url = await response.text();
+        return url;
+    }
+    async function getImageFile(div) {
+        let chatbot = document.getElementById("chatbot");
+        chatbot.style.height = "";
+        return new Promise((resolve, reject) =>
+            html2canvas(div)
+                .then((canvas) => {
+                    chatbot.style.height = "400px";
+                    const imageBlob = canvas.toBlob((blob) => {
+                        const imageId = Date.now();
+                        const fileName = "GILL-" + imageId + ".jpg";
+                        resolve(new File([blob], fileName, { type: 'image/jpeg' }));
+                    }, 'image/jpeg', 0.95);
+                })
+
+        )
+    }
+
+    const gradioEl = document.querySelector("gradio-app").shadowRoot || document.querySelector('body > gradio-app');
+    const chatbotEl = gradioEl.querySelector('#chatbot')
+    const imageFile = await getImageFile(chatbotEl);
+    console.log(imageFile);
+    const urlChatbotImage = await uploadFile(imageFile);
+    console.log(urlChatbotImage);
+    let titleTxt = `GILL Example`;
+
+    //const shareBtnEl = gradioEl.querySelector('#share-btn');
+    //shareBtnEl.style.pointerEvents = 'none';
+    const descriptionMd = `
+
+<img src='${urlChatbotImage}'>
+`;
+    const params = new URLSearchParams({
+        title: titleTxt,
+        description: descriptionMd,
+    });
+    const paramsStr = params.toString();
+    window.open(`https://huggingface.co/spaces/jykoh/gill/discussions/new?${paramsStr}`, '_blank');
+    //shareBtnEl.style.removeProperty('pointer-events');
+}
+"""
+
+save_js = """
+async () => {
+    const html2canvas = (await import('https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.esm.js')).default;
+
+    function saveAs(uri, filename) {
+      var link = document.createElement('a');
+      if (typeof link.download === 'string') {
+        link.href = uri;
+        link.download = filename;
+
+        //Firefox requires the link to be in the body
+        document.body.appendChild(link);
+
+        //simulate click
+        link.click();
+
+        //remove the link when done
+        document.body.removeChild(link);
+      } else {
+        window.open(uri);
+      }
+    }
+
+    async function getImageFile(div) {
+        let chatbot = document.getElementById("chatbot");
+        chatbot.style.height = "";
+        return new Promise((resolve, reject) =>
+            html2canvas(div)
+                .then((canvas) => {
+                    chatbot.style.height = "400px";
+                    const imageId = Date.now();
+                    const fileName = "GILL-" + imageId + ".png";
+                    saveAs(canvas.toDataURL(), fileName);
+                })
+
+        )
+    }
+    const gradioEl = document.querySelector("gradio-app").shadowRoot || document.querySelector('body > gradio-app');
+    const chatbotEl = gradioEl.querySelector('#chatbot')
+    const imageFile = await getImageFile(chatbotEl);
+    console.log(imageFile);
+}
+"""

inference_mmu.py

@@ -0,0 +1,174 @@
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+from PIL import Image
+from tqdm import tqdm
+import numpy as np
+import torch
+import wandb
+from models import Showo, MAGVITv2
+from prompting_utils import UniversalPrompting, create_attention_mask_for_mmu, create_attention_mask_for_mmu_vit
+from training.utils import get_config, flatten_omega_conf, image_transform
+from transformers import AutoTokenizer
+from models.clip_encoder import CLIPVisionTower
+from transformers import CLIPImageProcessor
+
+# import.training.conversation as conversation_lib
+from training import conversation as conversation_lib
+
+conversation_lib.default_conversation = conversation_lib.conv_templates["phi1.5"]
+SYSTEM_PROMPT = "A chat between a curious user and an artificial intelligence assistant. " \
+                "The assistant gives helpful, detailed, and polite answers to the user's questions."
+SYSTEM_PROMPT_LEN = 28
+
+def get_vq_model_class(model_type):
+    if model_type == "magvitv2":
+        return MAGVITv2
+    else:
+        raise ValueError(f"model_type {model_type} not supported.")
+
+if __name__ == '__main__':
+
+    config = get_config()
+
+    resume_wandb_run = config.wandb.resume
+    run_id = config.wandb.get("run_id", None)
+    if run_id is None:
+        resume_wandb_run = False
+        run_id = wandb.util.generate_id()
+        config.wandb.run_id = run_id
+
+    wandb_config = {k: v for k, v in flatten_omega_conf(config, resolve=True)}
+
+    wandb.init(
+        project="demo",
+        name=config.experiment.name + '_mmu',
+        config=wandb_config,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+    uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                       special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                       ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+    vq_model = get_vq_model_class(config.model.vq_model.type)
+    vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+    vq_model.requires_grad_(False)
+    vq_model.eval()
+
+    vision_tower_name = "openai/clip-vit-large-patch14-336"
+    vision_tower =  CLIPVisionTower(vision_tower_name).to(device)
+    clip_image_processor = CLIPImageProcessor.from_pretrained(vision_tower_name)
+
+    model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+    model.eval()
+
+    temperature = 0.8  # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
+    top_k = 1  # retain only the top_k most likely tokens, clamp others to have 0 probability
+
+    file_list = os.listdir(config.mmu_image_root)
+    responses = ['' for i in range(len(file_list))]
+    images = []
+    config.question = config.question.split(' *** ')
+    for i, file_name in enumerate(tqdm(file_list)):
+        image_path = os.path.join(config.mmu_image_root, file_name)
+        image_ori = Image.open(image_path).convert("RGB")
+        image = image_transform(image_ori, resolution=config.dataset.params.resolution).to(device)
+        image = image.unsqueeze(0)
+        images.append(image)
+
+        pixel_values = clip_image_processor.preprocess(image_ori, return_tensors="pt")["pixel_values"][0]
+
+        image_tokens = vq_model.get_code(image) + len(uni_prompting.text_tokenizer)
+        batch_size = 1
+
+        for question in config.question:
+            if config.model.showo.w_clip_vit:
+                conv = conversation_lib.default_conversation.copy()
+                conv.append_message(conv.roles[0], question)
+                conv.append_message(conv.roles[1], None)
+                prompt_question = conv.get_prompt()
+                question_input = []
+                question_input.append(prompt_question.strip())
+
+                input_ids_system = [uni_prompting.text_tokenizer(SYSTEM_PROMPT, return_tensors="pt", padding="longest").input_ids
+                                        for _ in range(batch_size)]
+                input_ids_system = torch.stack(input_ids_system, dim=0)
+                assert input_ids_system.shape[-1] == 28
+                input_ids_system = input_ids_system.to(device)
+                input_ids_system = input_ids_system[0]
+
+                input_ids = [uni_prompting.text_tokenizer(prompt, return_tensors="pt", padding="longest").input_ids
+                                for prompt in question_input]
+
+                input_ids = torch.stack(input_ids)
+                input_ids = torch.nn.utils.rnn.pad_sequence(
+                        input_ids, batch_first=True, padding_value=uni_prompting.text_tokenizer.pad_token_id
+                )
+                input_ids = torch.tensor(input_ids).to(device).squeeze(0)
+                # import pdb; pdb.set_trace()
+                input_ids_llava = torch.cat([
+                        (torch.ones(input_ids.shape[0], 1) *uni_prompting.sptids_dict['<|mmu|>']).to(device),
+                        input_ids_system,
+                        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
+                        # place your img embedding here
+                        (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
+                        input_ids,
+                ], dim=1).long()
+
+                images_embeddings = vision_tower(pixel_values[None])
+                images_embeddings = model.mm_projector(images_embeddings)
+
+                text_embeddings = model.showo.model.embed_tokens(input_ids_llava)
+
+                # Full input seq
+                part1 = text_embeddings[:, :2 + SYSTEM_PROMPT_LEN, :]
+                part2 = text_embeddings[:, 2 + SYSTEM_PROMPT_LEN:, :]
+                input_embeddings = torch.cat((part1, images_embeddings, part2), dim=1)
+
+                attention_mask_llava = create_attention_mask_for_mmu_vit(input_embeddings,
+                                                                        system_prompt_len=SYSTEM_PROMPT_LEN)
+
+                cont_toks_list = model.mmu_generate(input_embeddings=input_embeddings,
+                                                    attention_mask=attention_mask_llava[0].unsqueeze(0),
+                                                    max_new_tokens=100,
+                                                    top_k=top_k,
+                                                    eot_token=uni_prompting.sptids_dict['<|eot|>']
+                                                    )
+            else:
+                input_ids = uni_prompting.text_tokenizer(['USER: \n' + question + ' ASSISTANT:'])[
+                    'input_ids']
+                input_ids = torch.tensor(input_ids).to(device)
+
+                input_ids = torch.cat([
+                    (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|mmu|>']).to(device),
+                    (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|soi|>']).to(device),
+                    image_tokens,
+                    (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|eoi|>']).to(device),
+                    (torch.ones(input_ids.shape[0], 1) * uni_prompting.sptids_dict['<|sot|>']).to(device),
+                    input_ids
+                ], dim=1).long()
+
+                attention_mask = create_attention_mask_for_mmu(input_ids.to(device),
+                                                               eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']))
+
+                cont_toks_list = model.mmu_generate(input_ids, attention_mask=attention_mask,
+                                            max_new_tokens=100, top_k=top_k,
+                                            eot_token=uni_prompting.sptids_dict['<|eot|>'])
+
+            cont_toks_list = torch.stack(cont_toks_list).squeeze()[None]
+
+            text = uni_prompting.text_tokenizer.batch_decode(cont_toks_list, skip_special_tokens=True)
+            print(text)
+            responses[i] += f'User: ' + question + f'\n Answer : ' + text[0] + '\n'
+
+    images = torch.cat(images, dim=0)
+    images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+    images *= 255.0
+    images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+    pil_images = [Image.fromarray(image) for image in images]
+
+    wandb_images = [wandb.Image(image, caption=responses[i]) for i, image in enumerate(pil_images)]
+    wandb.log({"multimodal understanding": wandb_images}, step=0)
+

inference_t2i.py

@@ -0,0 +1,331 @@
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+from PIL import Image
+from tqdm import tqdm
+import numpy as np
+import torch
+import wandb
+from models import Showo, MAGVITv2, get_mask_chedule
+from prompting_utils import UniversalPrompting, create_attention_mask_predict_next
+from training.utils import get_config, flatten_omega_conf, image_transform
+from transformers import AutoTokenizer
+import torch.nn.functional as F
+
+def get_vq_model_class(model_type):
+    if model_type == "magvitv2":
+        return MAGVITv2
+    else:
+        raise ValueError(f"model_type {model_type} not supported.")
+
+if __name__ == '__main__':
+
+    config = get_config()
+
+    resume_wandb_run = config.wandb.resume
+    run_id = config.wandb.get("run_id", None)
+    if run_id is None:
+        resume_wandb_run = False
+        run_id = wandb.util.generate_id()
+        config.wandb.run_id = run_id
+
+    wandb_config = {k: v for k, v in flatten_omega_conf(config, resolve=True)}
+
+    wandb.init(
+        project="demo",
+        name=config.experiment.name + '_t2i' + f'_{config.mode}',
+        config=wandb_config,
+    )
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(config.model.showo.llm_model_path, padding_side="left")
+
+    uni_prompting = UniversalPrompting(tokenizer, max_text_len=config.dataset.preprocessing.max_seq_length,
+                                       special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                                       ignore_id=-100, cond_dropout_prob=config.training.cond_dropout_prob)
+
+    vq_model = get_vq_model_class(config.model.vq_model.type)
+    vq_model = vq_model.from_pretrained(config.model.vq_model.vq_model_name).to(device)
+    vq_model.requires_grad_(False)
+    vq_model.eval()
+
+    model = Showo.from_pretrained(config.model.showo.pretrained_model_path).to(device)
+    model.eval()
+
+    mask_token_id = model.config.mask_token_id
+
+    # load from users passed arguments
+    if config.get("validation_prompts_file", None) is not None:
+        config.dataset.params.validation_prompts_file = config.validation_prompts_file
+    config.training.batch_size = config.batch_size
+    config.training.guidance_scale = config.guidance_scale
+    config.training.generation_timesteps = config.generation_timesteps
+    # load from users passed arguments
+
+    if config.mode == 'inpainting':
+
+        prompt = [config.prompt] * config.batch_size
+        inpainting_image = Image.open(config.image_path).convert("RGB")
+        inpainting_mask = Image.open(config.inpainting_mask_path).convert("L")
+
+        import pdb
+        pdb.set_trace()
+
+        inpainting_image = image_transform(inpainting_image, resolution=config.dataset.params.resolution).to(device)
+        inpainting_mask = image_transform(inpainting_mask, resolution=config.dataset.params.resolution, normalize=False)
+
+        # record original image and inpainting mask
+        images = torch.clamp(
+            (torch.stack([inpainting_image, inpainting_mask.repeat(3, 1, 1).to(device)], dim=0) + 1.0) / 2.0,
+            min=0.0, max=1.0)
+        images *= 255.0
+        images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+        pil_images = [Image.fromarray(image) for image in images]
+
+        labels = ['original image', 'inpainting mask']
+        wandb_images = [wandb.Image(image, caption=labels[i]) for i, image in enumerate(pil_images)]
+
+        inpainting_image = inpainting_image.unsqueeze(0).repeat(config.training.batch_size, 1, 1, 1)
+
+        inpainting_mask = inpainting_mask.unsqueeze(0).to(device)
+        inpainting_mask = F.interpolate(inpainting_mask, size=config.dataset.params.resolution // 16, mode='bicubic')
+        inpainting_mask = inpainting_mask.repeat(config.training.batch_size, 1, 1, 1)
+
+        inpainting_mask[inpainting_mask < 0.5] = 0
+        inpainting_mask[inpainting_mask >= 0.5] = 1
+
+        inpainting_mask = inpainting_mask.reshape(config.training.batch_size, -1)
+        inpainting_mask = inpainting_mask.to(torch.bool)
+
+        inpainting_image_tokens = vq_model.get_code(inpainting_image) + len(uni_prompting.text_tokenizer)
+        inpainting_image_tokens[inpainting_mask] = mask_token_id
+
+        input_ids, _ = uni_prompting((prompt, inpainting_image_tokens), 't2i_gen')
+
+        if config.training.guidance_scale > 0:
+            uncond_input_ids, _ = uni_prompting(([''] * len(prompt), inpainting_image_tokens), 't2i_gen')
+            attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+        else:
+            attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                rm_pad_in_image=True)
+            uncond_input_ids = None
+
+        if config.get("mask_schedule", None) is not None:
+            schedule = config.mask_schedule.schedule
+            args = config.mask_schedule.get("params", {})
+            mask_schedule = get_mask_chedule(schedule, **args)
+        else:
+            mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+        with torch.no_grad():
+            gen_token_ids = model.t2i_generate(
+                input_ids=input_ids,
+                uncond_input_ids=uncond_input_ids,
+                attention_mask=attention_mask,
+                guidance_scale=config.training.guidance_scale,
+                temperature=config.training.get("generation_temperature", 1.0),
+                timesteps=config.training.generation_timesteps,
+                noise_schedule=mask_schedule,
+                noise_type=config.training.get("noise_type", "mask"),
+                seq_len=config.model.showo.num_vq_tokens,
+                uni_prompting=uni_prompting,
+                config=config,
+            )
+
+        gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+        images = vq_model.decode_code(gen_token_ids)
+
+        images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+        images *= 255.0
+        images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+        pil_images = [Image.fromarray(image) for image in images]
+        # import ipdb
+        # ipdb.set_trace()
+        wandb_images.extend([wandb.Image(image, caption=prompt[i]) for i, image in enumerate(pil_images)])
+        wandb.log({"generated_images": wandb_images}, step=0)
+
+    elif config.mode == 'extrapolation':
+
+        prompt = [p for p in config.prompt.split(" *** ") if len(p) != 0]
+        extra_direction = [d for d in config.extra_direction.split(" *** ") if len(d) != 0]
+        print(prompt, extra_direction)
+        W = config.dataset.params.resolution // 16
+        for id, (prt, direction) in enumerate(zip(prompt, extra_direction)):
+            prt = [prt] * config.training.batch_size
+            if id == 0:
+                extrapolation_image = Image.open(config.image_path).convert("RGB")
+                extrapolation_image = image_transform(extrapolation_image,
+                                                      resolution=config.dataset.params.resolution).to(device)
+
+                B, _, _ = extrapolation_image.shape
+                extrapolation_image = extrapolation_image.unsqueeze(0)
+                extrapolation_image_tokens = vq_model.get_code(extrapolation_image) + len(uni_prompting.text_tokenizer)
+                extrapolation_image_tokens = extrapolation_image_tokens.reshape(1,
+                                                                                config.dataset.params.resolution // 16,
+                                                                                config.dataset.params.resolution // 16)
+                extrapolation_image_tokens = extrapolation_image_tokens.repeat(config.training.batch_size, 1, 1)
+            else:
+
+
+                extrapolation_image_tokens = gen_token_ids + len(uni_prompting.text_tokenizer)
+
+            image_left_part = extrapolation_image_tokens[:, :, :-(W//2-config.offset)] - len(uni_prompting.text_tokenizer)
+            image_right_part = extrapolation_image_tokens[:, :, W//2-config.offset:] - len(uni_prompting.text_tokenizer)
+            image_up_part = extrapolation_image_tokens[:, :-(W//2-config.offset), :] - len(uni_prompting.text_tokenizer)
+            image_down_part = extrapolation_image_tokens[:, W//2-config.offset:, :] - len(uni_prompting.text_tokenizer)
+
+            if direction in ['left', 'right']:
+                extrapolation_mask = torch.zeros((config.training.batch_size,
+                                                  config.dataset.params.resolution // 16,
+                                                  config.dataset.params.resolution // 16 // 2 + config.offset),
+                                                 dtype=torch.int64, device=device) + mask_token_id
+            else:
+                extrapolation_mask = torch.zeros((config.training.batch_size,
+                                                  config.dataset.params.resolution // 16 // 2 + config.offset,
+                                                  config.dataset.params.resolution // 16),
+                                                 dtype=torch.int64, device=device) + mask_token_id
+
+            if direction == 'left':
+                extrapolation_image_tokens = torch.cat(
+                    [extrapolation_mask, extrapolation_image_tokens[:, :, :W//2-config.offset]], dim=-1)
+            elif direction == 'right':
+                extrapolation_image_tokens = torch.cat(
+                    [extrapolation_image_tokens[:, :, -(W//2-config.offset):], extrapolation_mask], dim=-1)
+            elif direction == 'up':
+                extrapolation_image_tokens = torch.cat(
+                    [extrapolation_mask, extrapolation_image_tokens[:, :W // 2 - config.offset, :]], dim=-2)
+            else:
+                extrapolation_image_tokens = torch.cat(
+                    [extrapolation_image_tokens[:, -(W // 2 - config.offset):, :], extrapolation_mask], dim=-2)
+
+            extrapolation_image_tokens = extrapolation_image_tokens.reshape(config.training.batch_size, -1)
+
+            input_ids, _ = uni_prompting((prt, extrapolation_image_tokens), 't2i_gen')
+
+            if config.training.guidance_scale > 0:
+                uncond_input_ids, _ = uni_prompting(([''] * len(prt), extrapolation_image_tokens), 't2i_gen')
+                attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                    pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                    soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                    eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                    rm_pad_in_image=True)
+            else:
+                attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                    pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                    soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                    eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                    rm_pad_in_image=True)
+                uncond_input_ids = None
+
+            if config.get("mask_schedule", None) is not None:
+                schedule = config.mask_schedule.schedule
+                args = config.mask_schedule.get("params", {})
+                mask_schedule = get_mask_chedule(schedule, **args)
+            else:
+                mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+            with torch.no_grad():
+                gen_token_ids = model.t2i_generate(
+                    input_ids=input_ids,
+                    uncond_input_ids=uncond_input_ids,
+                    attention_mask=attention_mask,
+                    guidance_scale=config.training.guidance_scale,
+                    temperature=config.training.get("generation_temperature", 1.0),
+                    timesteps=config.training.generation_timesteps,
+                    noise_schedule=mask_schedule,
+                    noise_type=config.training.get("noise_type", "mask"),
+                    seq_len=config.model.showo.num_vq_tokens,
+                    uni_prompting=uni_prompting,
+                    config=config,
+                )
+
+            gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+            gen_token_ids = gen_token_ids.reshape(config.training.batch_size,
+                                                  config.dataset.params.resolution // 16,
+                                                  config.dataset.params.resolution // 16)
+            if direction == 'left':
+                gen_token_ids = torch.cat([gen_token_ids, image_right_part], dim=-1)
+            elif direction == 'right':
+                gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-1)
+            elif direction == 'up':
+                gen_token_ids = torch.cat([gen_token_ids, image_down_part], dim=-2)
+            else:
+                gen_token_ids = torch.cat([image_left_part, gen_token_ids], dim=-2)
+
+        _, h, w = gen_token_ids.shape
+        gen_token_ids = gen_token_ids.reshape(config.training.batch_size, -1)
+        images = vq_model.decode_code(gen_token_ids, shape=(h, w))
+
+        images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+        images *= 255.0
+        images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+        pil_images = [Image.fromarray(image) for image in images]
+
+        wandb_images = [wandb.Image(image, caption=' '.join(prompt)) for i, image in enumerate(pil_images)]
+        wandb.log({"generated_images": wandb_images}, step=0)
+
+    elif config.mode == 't2i':
+        with open(config.dataset.params.validation_prompts_file, "r") as f:
+            validation_prompts = f.read().splitlines()
+
+        for step in tqdm(range(0, len(validation_prompts), config.training.batch_size)):
+            prompts = validation_prompts[step:step + config.training.batch_size]
+
+            image_tokens = torch.ones((len(prompts), config.model.showo.num_vq_tokens),
+                                      dtype=torch.long, device=device) * mask_token_id
+
+            input_ids, _ = uni_prompting((prompts, image_tokens), 't2i_gen')
+
+            if config.training.guidance_scale > 0:
+                uncond_input_ids, _ = uni_prompting(([''] * len(prompts), image_tokens), 't2i_gen')
+                attention_mask = create_attention_mask_predict_next(torch.cat([input_ids, uncond_input_ids], dim=0),
+                                                                    pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                    soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                    eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                    rm_pad_in_image=True)
+            else:
+                attention_mask = create_attention_mask_predict_next(input_ids,
+                                                                    pad_id=int(uni_prompting.sptids_dict['<|pad|>']),
+                                                                    soi_id=int(uni_prompting.sptids_dict['<|soi|>']),
+                                                                    eoi_id=int(uni_prompting.sptids_dict['<|eoi|>']),
+                                                                    rm_pad_in_image=True)
+                uncond_input_ids = None
+
+            if config.get("mask_schedule", None) is not None:
+                schedule = config.mask_schedule.schedule
+                args = config.mask_schedule.get("params", {})
+                mask_schedule = get_mask_chedule(schedule, **args)
+            else:
+                mask_schedule = get_mask_chedule(config.training.get("mask_schedule", "cosine"))
+
+            with torch.no_grad():
+                gen_token_ids = model.t2i_generate(
+                    input_ids=input_ids,
+                    uncond_input_ids=uncond_input_ids,
+                    attention_mask=attention_mask,
+                    guidance_scale=config.training.guidance_scale,
+                    temperature=config.training.get("generation_temperature", 1.0),
+                    timesteps=config.training.generation_timesteps,
+                    noise_schedule=mask_schedule,
+                    noise_type=config.training.get("noise_type", "mask"),
+                    seq_len=config.model.showo.num_vq_tokens,
+                    uni_prompting=uni_prompting,
+                    config=config,
+                )
+
+            gen_token_ids = torch.clamp(gen_token_ids, max=config.model.showo.codebook_size - 1, min=0)
+            images = vq_model.decode_code(gen_token_ids)
+
+            images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
+            images *= 255.0
+            images = images.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+            pil_images = [Image.fromarray(image) for image in images]
+
+            wandb_images = [wandb.Image(image, caption=prompts[i]) for i, image in enumerate(pil_images)]
+            wandb.log({"generated_images": wandb_images}, step=step)

models/__init__.py

@@ -0,0 +1,4 @@
+from .modeling_showo import Showo
+from .modeling_magvitv2 import VQGANEncoder, VQGANDecoder, LFQuantizer, MAGVITv2
+from .sampling import *
+from .clip_encoder import CLIPVisionTower

models/clip_encoder.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+
+from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+
+class CLIPVisionTower(nn.Module):
+    def __init__(self, vision_tower):
+        super().__init__()
+
+        self.is_loaded = False
+
+        self.vision_tower_name = vision_tower
+        self.select_layer = -2
+        self.select_feature = "patch"
+        self.load_model()
+        self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
+
+    def load_model(self, device_map=None):
+        if self.is_loaded:
+            print('{} is already loaded, `load_model` called again, skipping.'.format(self.vision_tower_name))
+            return
+
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name, device_map=device_map)
+        self.vision_tower.requires_grad_(False)
+
+        self.is_loaded = True
+
+    def feature_select(self, image_forward_outs):
+        image_features = image_forward_outs.hidden_states[self.select_layer]
+        if self.select_feature == 'patch':
+            image_features = image_features[:, 1:]
+        elif self.select_feature == 'cls_patch':
+            image_features = image_features
+        else:
+            raise ValueError(f'Unexpected select feature: {self.select_feature}')
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_forward_out = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
+                image_feature = self.feature_select(image_forward_out).to(image.dtype)
+                image_features.append(image_feature)
+        else:
+            image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
+            image_features = self.feature_select(image_forward_outs).to(images.dtype)
+
+        return image_features
+
+    @property
+    def dummy_feature(self):
+        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device
+
+    @property
+    def config(self):
+        if self.is_loaded:
+            return self.vision_tower.config
+        else:
+            return self.cfg_only
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches_per_side(self):
+        return self.config.image_size // self.config.patch_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+class CLIPVisionTowerS2(CLIPVisionTower):
+    def __init__(self, vision_tower, args, delay_load=False):
+        super().__init__(vision_tower, args, delay_load)
+
+        self.s2_scales = getattr(args, 's2_scales', '336,672,1008')
+        self.s2_scales = list(map(int, self.s2_scales.split(',')))
+        self.s2_scales.sort()
+        self.s2_split_size = self.s2_scales[0]
+        self.s2_image_size = self.s2_scales[-1]
+
+        try:
+            from s2wrapper import forward as multiscale_forward
+        except ImportError:
+            raise ImportError('Package s2wrapper not found! Please install by running: \npip install git+https://github.com/bfshi/scaling_on_scales.git')
+        self.multiscale_forward = multiscale_forward
+
+        # change resize/crop size in preprocessing to the largest image size in s2_scale
+        if not delay_load or getattr(args, 'unfreeze_mm_vision_tower', False):
+            self.image_processor.size['shortest_edge'] = self.s2_image_size
+            self.image_processor.crop_size['height'] = self.image_processor.crop_size['width'] = self.s2_image_size
+
+    def load_model(self, device_map=None):
+        if self.is_loaded:
+            print('{} is already loaded, `load_model` called again, skipping.'.format(self.vision_tower_name))
+            return
+
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+        self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name, device_map=device_map)
+        self.vision_tower.requires_grad_(False)
+
+        self.image_processor.size['shortest_edge'] = self.s2_image_size
+        self.image_processor.crop_size['height'] = self.image_processor.crop_size['width'] = self.s2_image_size
+
+        self.is_loaded = True
+
+    @torch.no_grad()
+    def forward_feature(self, images):
+        image_forward_outs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
+        image_features = self.feature_select(image_forward_outs).to(images.dtype)
+        return image_features
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            for image in images:
+                image_feature = self.multiscale_forward(self.forward_feature, image.unsqueeze(0), img_sizes=self.s2_scales, max_split_size=self.s2_split_size)
+                image_features.append(image_feature)
+        else:
+            image_features = self.multiscale_forward(self.forward_feature, images, img_sizes=self.s2_scales, max_split_size=self.s2_split_size)
+
+        return image_features
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size * len(self.s2_scales)

models/common_modules.py

@@ -0,0 +1,407 @@
+"""
+Modified from https://github.com/CompVis/taming-transformers/blob/master/taming/modules/diffusionmodules/model.py#L34
+"""
+
+import math
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+
+def nonlinearity(x):
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(
+        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class DepthToSpaceUpsample(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+    ):
+        super().__init__()
+        conv = nn.Conv2d(in_channels, in_channels * 4, 1)
+
+        self.net = nn.Sequential(
+            conv,
+            nn.SiLU(),
+            Rearrange("b (c p1 p2) h w -> b c (h p1) (w p2)", p1=2, p2=2),
+        )
+
+        self.init_conv_(conv)
+
+    def init_conv_(self, conv):
+        o, i, h, w = conv.weight.shape
+        conv_weight = torch.empty(o // 4, i, h, w)
+        nn.init.kaiming_uniform_(conv_weight)
+        conv_weight = repeat(conv_weight, "o ... -> (o 4) ...")
+
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def forward(self, x):
+        out = self.net(x)
+        return out
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+def unpack_time(t, batch):
+    _, c, w, h = t.size()
+    out = torch.reshape(t, [batch, -1, c, w, h])
+    out = rearrange(out, "b t c h w -> b c t h w")
+    return out
+
+
+def pack_time(t):
+    out = rearrange(t, "b c t h w -> b t c h w")
+    _, _, c, w, h = out.size()
+    return torch.reshape(out, [-1, c, w, h])
+
+
+class TimeDownsample2x(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out=None,
+        kernel_size=3,
+    ):
+        super().__init__()
+        if dim_out is None:
+            dim_out = dim
+        self.time_causal_padding = (kernel_size - 1, 0)
+        self.conv = nn.Conv1d(dim, dim_out, kernel_size, stride=2)
+
+    def forward(self, x):
+        x = rearrange(x, "b c t h w -> b h w c t")
+        b, h, w, c, t = x.size()
+        x = torch.reshape(x, [-1, c, t])
+
+        x = F.pad(x, self.time_causal_padding)
+        out = self.conv(x)
+
+        out = torch.reshape(out, [b, h, w, c, t])
+        out = rearrange(out, "b h w c t -> b c t h w")
+        out = rearrange(out, "b h w c t -> b c t h w")
+        return out
+
+
+class TimeUpsample2x(nn.Module):
+    def __init__(self, dim, dim_out=None):
+        super().__init__()
+        if dim_out is None:
+            dim_out = dim
+        conv = nn.Conv1d(dim, dim_out * 2, 1)
+
+        self.net = nn.Sequential(
+            nn.SiLU(), conv, Rearrange("b (c p) t -> b c (t p)", p=2)
+        )
+
+        self.init_conv_(conv)
+
+    def init_conv_(self, conv):
+        o, i, t = conv.weight.shape
+        conv_weight = torch.empty(o // 2, i, t)
+        nn.init.kaiming_uniform_(conv_weight)
+        conv_weight = repeat(conv_weight, "o ... -> (o 2) ...")
+
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def forward(self, x):
+        x = rearrange(x, "b c t h w -> b h w c t")
+        b, h, w, c, t = x.size()
+        x = torch.reshape(x, [-1, c, t])
+
+        out = self.net(x)
+        out = out[:, :, 1:].contiguous()
+
+        out = torch.reshape(out, [b, h, w, c, t])
+        out = rearrange(out, "b h w c t -> b c t h w")
+        return out
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.k = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)  # b,hw,c
+        k = k.reshape(b, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class TimeAttention(AttnBlock):
+    def forward(self, x, *args, **kwargs):
+        x = rearrange(x, "b c t h w -> b h w t c")
+        b, h, w, t, c = x.size()
+        x = torch.reshape(x, (-1, t, c))
+
+        x = super().forward(x, *args, **kwargs)
+
+        x = torch.reshape(x, [b, h, w, t, c])
+        return rearrange(x, "b h w t c -> b c t h w")
+
+
+class Residual(nn.Module):
+    def __init__(self, fn: nn.Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+
+def cast_tuple(t, length=1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+
+class CausalConv3d(nn.Module):
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        pad_mode="constant",
+        **kwargs
+    ):
+        super().__init__()
+        kernel_size = cast_tuple(kernel_size, 3)
+
+        time_kernel_size, height_kernel_size, width_kernel_size = kernel_size
+
+        dilation = kwargs.pop("dilation", 1)
+        stride = kwargs.pop("stride", 1)
+
+        self.pad_mode = pad_mode
+        time_pad = dilation * (time_kernel_size - 1) + (1 - stride)
+        height_pad = height_kernel_size // 2
+        width_pad = width_kernel_size // 2
+
+        self.time_pad = time_pad
+        self.time_causal_padding = (
+            width_pad,
+            width_pad,
+            height_pad,
+            height_pad,
+            time_pad,
+            0,
+        )
+
+        stride = (stride, 1, 1)
+        dilation = (dilation, 1, 1)
+        self.conv = nn.Conv3d(
+            chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs
+        )
+
+    def forward(self, x):
+        pad_mode = self.pad_mode if self.time_pad < x.shape[2] else "constant"
+
+        x = F.pad(x, self.time_causal_padding, mode=pad_mode)
+        return self.conv(x)
+
+
+def ResnetBlockCausal3D(
+    dim, kernel_size: Union[int, Tuple[int, int, int]], pad_mode: str = "constant"
+):
+    net = nn.Sequential(
+        Normalize(dim),
+        nn.SiLU(),
+        CausalConv3d(dim, dim, kernel_size, pad_mode),
+        Normalize(dim),
+        nn.SiLU(),
+        CausalConv3d(dim, dim, kernel_size, pad_mode),
+    )
+    return Residual(net)
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout,
+        temb_channels=512
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+        else:
+            self.temb_proj = None
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class DinoV2Model(nn.Module):
+    def __init__(
+        self,
+        model_name,
+        local_checkpoint_path="",
+        renorm_input=False,
+        old_input_mean=0.5,
+        old_input_std=0.5,
+        freeze_model=False,
+    ):
+        super().__init__()
+        if local_checkpoint_path != "":
+            self._model = torch.hub.load(
+                local_checkpoint_path, model_name, source="local"
+            )
+        else:
+            self._model = torch.hub.load("facebookresearch/dinov2", model_name)
+        self.register_buffer(
+            "_dino_input_mean",
+            torch.tensor([0.485, 0.456, 0.406]).float()[None, :, None, None],
+        )
+        self.register_buffer(
+            "_dino_input_std",
+            torch.tensor([0.229, 0.224, 0.225]).float()[None, :, None, None],
+        )
+        self._old_input_mean = old_input_mean
+        self._old_input_std = old_input_std
+        self._renorm_input = renorm_input
+        if freeze_model:
+            for param in self._model.parameters():
+                param.requires_grad = False
+
+    def forward(self, inputs):
+        batch, _, height, width = inputs.size()
+        if self._renorm_input:
+            inputs = inputs * self._old_input_mean + self._old_input_std
+            inputs = (inputs - self._dino_input_mean) / self._dino_input_std
+        # TODO(yanwan): If we want to remove this resizing, have to modify the decoder to support upscaling by a factor of 14.
+        # Reduce both height and width to 7/8 of their original values while maintaining aspect ratio to fit dinov2 requirement.
+        new_height = height // 8 * 7
+        new_width = width // 8 * 7
+        inputs = F.interpolate(inputs, (new_height, new_width), mode="bilinear")
+        features = self._model.forward_features(inputs)["x_norm_patchtokens"]
+        features = torch.transpose(features, 1, 2).contiguous()
+        features = torch.reshape(
+            features, (batch, -1, new_height // 14, new_width // 14)
+        )
+        return features

models/logging.py

@@ -0,0 +1,338 @@
+# coding=utf-8
+# Copyright 2023 Optuna, Hugging Face
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Logging utilities."""
+
+import logging
+import os
+import sys
+import threading
+from logging import CRITICAL  # NOQA
+from logging import DEBUG  # NOQA
+from logging import ERROR  # NOQA
+from logging import FATAL  # NOQA
+from logging import INFO  # NOQA
+from logging import NOTSET  # NOQA
+from logging import WARN  # NOQA
+from logging import WARNING  # NOQA
+from typing import Optional
+
+from tqdm import auto as tqdm_lib
+
+_lock = threading.Lock()
+_default_handler: Optional[logging.Handler] = None
+
+log_levels = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.WARNING
+
+_tqdm_active = True
+
+
+def _get_default_logging_level():
+    """
+    If muse_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+    not - fall back to `_default_log_level`
+    """
+    env_level_str = os.getenv("muse_VERBOSITY", None)
+    if env_level_str:
+        if env_level_str in log_levels:
+            return log_levels[env_level_str]
+        else:
+            logging.getLogger().warning(
+                f"Unknown option muse_VERBOSITY={env_level_str}, has to be one of: { ', '.join(log_levels.keys()) }"
+            )
+    return _default_log_level
+
+
+def _get_library_name() -> str:
+    return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+    return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if _default_handler:
+            # This library has already configured the library root logger.
+            return
+        _default_handler = logging.StreamHandler()  # Set sys.stderr as stream.
+        _default_handler.flush = sys.stderr.flush
+
+        # Apply our default configuration to the library root logger.
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.addHandler(_default_handler)
+        library_root_logger.setLevel(_get_default_logging_level())
+        library_root_logger.propagate = False
+
+
+def _reset_library_root_logger() -> None:
+    global _default_handler
+
+    with _lock:
+        if not _default_handler:
+            return
+
+        library_root_logger = _get_library_root_logger()
+        library_root_logger.removeHandler(_default_handler)
+        library_root_logger.setLevel(logging.NOTSET)
+        _default_handler = None
+
+
+def get_log_levels_dict():
+    return log_levels
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Return a logger with the specified name.
+
+    This function is not supposed to be directly accessed unless you are writing a custom muse module.
+    """
+
+    if name is None:
+        name = _get_library_name()
+
+    _configure_library_root_logger()
+    return logging.getLogger(name)
+
+
+def get_verbosity() -> int:
+    """
+    Return the current level for the 🤗 muse' root logger as an int.
+
+    Returns:
+        `int`: The logging level.
+
+    <Tip>
+
+    🤗 muse has following logging levels:
+
+    - 50: `muse.logging.CRITICAL` or `muse.logging.FATAL`
+    - 40: `muse.logging.ERROR`
+    - 30: `muse.logging.WARNING` or `muse.logging.WARN`
+    - 20: `muse.logging.INFO`
+    - 10: `muse.logging.DEBUG`
+
+    </Tip>"""
+
+    _configure_library_root_logger()
+    return _get_library_root_logger().getEffectiveLevel()
+
+
+def set_verbosity(verbosity: int) -> None:
+    """
+    Set the verbosity level for the 🤗 muse' root logger.
+
+    Args:
+        verbosity (`int`):
+            Logging level, e.g., one of:
+
+            - `muse.logging.CRITICAL` or `muse.logging.FATAL`
+            - `muse.logging.ERROR`
+            - `muse.logging.WARNING` or `muse.logging.WARN`
+            - `muse.logging.INFO`
+            - `muse.logging.DEBUG`
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().setLevel(verbosity)
+
+
+def set_verbosity_info():
+    """Set the verbosity to the `INFO` level."""
+    return set_verbosity(INFO)
+
+
+def set_verbosity_warning():
+    """Set the verbosity to the `WARNING` level."""
+    return set_verbosity(WARNING)
+
+
+def set_verbosity_debug():
+    """Set the verbosity to the `DEBUG` level."""
+    return set_verbosity(DEBUG)
+
+
+def set_verbosity_error():
+    """Set the verbosity to the `ERROR` level."""
+    return set_verbosity(ERROR)
+
+
+def disable_default_handler() -> None:
+    """Disable the default handler of the HuggingFace muse' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().removeHandler(_default_handler)
+
+
+def enable_default_handler() -> None:
+    """Enable the default handler of the HuggingFace muse' root logger."""
+
+    _configure_library_root_logger()
+
+    assert _default_handler is not None
+    _get_library_root_logger().addHandler(_default_handler)
+
+
+def add_handler(handler: logging.Handler) -> None:
+    """adds a handler to the HuggingFace muse' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None
+    _get_library_root_logger().addHandler(handler)
+
+
+def remove_handler(handler: logging.Handler) -> None:
+    """removes given handler from the HuggingFace muse' root logger."""
+
+    _configure_library_root_logger()
+
+    assert handler is not None and handler not in _get_library_root_logger().handlers
+    _get_library_root_logger().removeHandler(handler)
+
+
+def disable_propagation() -> None:
+    """
+    Disable propagation of the library log outputs. Note that log propagation is disabled by default.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = False
+
+
+def enable_propagation() -> None:
+    """
+    Enable propagation of the library log outputs. Please disable the HuggingFace muse' default handler to prevent
+    double logging if the root logger has been configured.
+    """
+
+    _configure_library_root_logger()
+    _get_library_root_logger().propagate = True
+
+
+def enable_explicit_format() -> None:
+    """
+    Enable explicit formatting for every HuggingFace muse' logger. The explicit formatter is as follows:
+    ```
+        [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE
+    ```
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s")
+        handler.setFormatter(formatter)
+
+
+def reset_format() -> None:
+    """
+    Resets the formatting for HuggingFace muse' loggers.
+
+    All handlers currently bound to the root logger are affected by this method.
+    """
+    handlers = _get_library_root_logger().handlers
+
+    for handler in handlers:
+        handler.setFormatter(None)
+
+
+def warning_advice(self, *args, **kwargs):
+    """
+    This method is identical to `logger.warning()`, but if env var muse_NO_ADVISORY_WARNINGS=1 is set, this
+    warning will not be printed
+    """
+    no_advisory_warnings = os.getenv("muse_NO_ADVISORY_WARNINGS", False)
+    if no_advisory_warnings:
+        return
+    self.warning(*args, **kwargs)
+
+
+logging.Logger.warning_advice = warning_advice
+
+
+class EmptyTqdm:
+    """Dummy tqdm which doesn't do anything."""
+
+    def __init__(self, *args, **kwargs):  # pylint: disable=unused-argument
+        self._iterator = args[0] if args else None
+
+    def __iter__(self):
+        return iter(self._iterator)
+
+    def __getattr__(self, _):
+        """Return empty function."""
+
+        def empty_fn(*args, **kwargs):  # pylint: disable=unused-argument
+            return
+
+        return empty_fn
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type_, value, traceback):
+        return
+
+
+class _tqdm_cls:
+    def __call__(self, *args, **kwargs):
+        if _tqdm_active:
+            return tqdm_lib.tqdm(*args, **kwargs)
+        else:
+            return EmptyTqdm(*args, **kwargs)
+
+    def set_lock(self, *args, **kwargs):
+        self._lock = None
+        if _tqdm_active:
+            return tqdm_lib.tqdm.set_lock(*args, **kwargs)
+
+    def get_lock(self):
+        if _tqdm_active:
+            return tqdm_lib.tqdm.get_lock()
+
+
+tqdm = _tqdm_cls()
+
+
+def is_progress_bar_enabled() -> bool:
+    """Return a boolean indicating whether tqdm progress bars are enabled."""
+    global _tqdm_active
+    return bool(_tqdm_active)
+
+
+def enable_progress_bar():
+    """Enable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = True
+
+
+def disable_progress_bar():
+    """Disable tqdm progress bar."""
+    global _tqdm_active
+    _tqdm_active = False

models/lr_schedulers.py

@@ -0,0 +1,292 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch optimization for diffusion models."""
+
+import math
+from enum import Enum
+from typing import Optional, Union
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+from .logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class SchedulerType(Enum):
+    LINEAR = "linear"
+    COSINE = "cosine"
+    COSINE_WITH_RESTARTS = "cosine_with_restarts"
+    POLYNOMIAL = "polynomial"
+    CONSTANT = "constant"
+    CONSTANT_WITH_WARMUP = "constant_with_warmup"
+
+
+def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
+    """
+    Create a schedule with a constant learning rate, using the learning rate set in optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+    return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
+
+
+def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
+    """
+    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
+    increases linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step: int):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1.0, num_warmup_steps))
+        return 1.0
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
+
+
+def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
+    """
+    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step: int):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        return max(
+            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
+        )
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_cosine_schedule_with_warmup(
+    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+):
+    """
+    Create a schedule with a learning rate that decreases following the values of the cosine function between the
+    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+    initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        num_periods (`float`, *optional*, defaults to 0.5):
+            The number of periods of the cosine function in a schedule (the default is to just decrease from the max
+            value to 0 following a half-cosine).
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_cosine_with_hard_restarts_schedule_with_warmup(
+    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
+):
+    """
+    Create a schedule with a learning rate that decreases following the values of the cosine function between the
+    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
+    linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        num_cycles (`int`, *optional*, defaults to 1):
+            The number of hard restarts to use.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        if progress >= 1.0:
+            return 0.0
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def get_polynomial_decay_schedule_with_warmup(
+    optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
+):
+    """
+    Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
+    optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the
+    initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        lr_end (`float`, *optional*, defaults to 1e-7):
+            The end LR.
+        power (`float`, *optional*, defaults to 1.0):
+            Power factor.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
+    implementation at
+    https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+
+    """
+
+    lr_init = optimizer.defaults["lr"]
+    if not (lr_init > lr_end):
+        raise ValueError(f"lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})")
+
+    def lr_lambda(current_step: int):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        elif current_step > num_training_steps:
+            return lr_end / lr_init  # as LambdaLR multiplies by lr_init
+        else:
+            lr_range = lr_init - lr_end
+            decay_steps = num_training_steps - num_warmup_steps
+            pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
+            decay = lr_range * pct_remaining**power + lr_end
+            return decay / lr_init  # as LambdaLR multiplies by lr_init
+
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+TYPE_TO_SCHEDULER_FUNCTION = {
+    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
+    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
+    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
+    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
+    SchedulerType.CONSTANT: get_constant_schedule,
+    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
+}
+
+
+def get_scheduler(
+    name: Union[str, SchedulerType],
+    optimizer: Optimizer,
+    num_warmup_steps: Optional[int] = None,
+    num_training_steps: Optional[int] = None,
+    num_cycles: int = 1,
+    power: float = 1.0,
+):
+    """
+    Unified API to get any scheduler from its name.
+
+    Args:
+        name (`str` or `SchedulerType`):
+            The name of the scheduler to use.
+        optimizer (`torch.optim.Optimizer`):
+            The optimizer that will be used during training.
+        num_warmup_steps (`int`, *optional*):
+            The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+            optional), the function will raise an error if it's unset and the scheduler type requires it.
+        num_training_steps (`int``, *optional*):
+            The number of training steps to do. This is not required by all schedulers (hence the argument being
+            optional), the function will raise an error if it's unset and the scheduler type requires it.
+        num_cycles (`int`, *optional*):
+            The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
+        power (`float`, *optional*, defaults to 1.0):
+            Power factor. See `POLYNOMIAL` scheduler
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+    """
+    name = SchedulerType(name)
+    schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+    if name == SchedulerType.CONSTANT:
+        return schedule_func(optimizer)
+
+    # All other schedulers require `num_warmup_steps`
+    if num_warmup_steps is None:
+        raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+    if name == SchedulerType.CONSTANT_WITH_WARMUP:
+        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+    # All other schedulers require `num_training_steps`
+    if num_training_steps is None:
+        raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+    if name == SchedulerType.COSINE_WITH_RESTARTS:
+        return schedule_func(
+            optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
+        )
+
+    if name == SchedulerType.POLYNOMIAL:
+        return schedule_func(
+            optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
+        )
+
+    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)

models/misc.py

@@ -0,0 +1,53 @@
+from omegaconf import OmegaConf
+import torch
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    NamedTuple,
+    NewType,
+    Optional,
+    Sized,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+# Tensor dtype
+# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
+from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
+
+# Config type
+from omegaconf import DictConfig
+
+# PyTorch Tensor type
+from torch import Tensor
+
+# Runtime type checking decorator
+from typeguard import typechecked as typechecker
+
+
+def broadcast(tensor, src=0):
+    if not _distributed_available():
+        return tensor
+    else:
+        torch.distributed.broadcast(tensor, src=src)
+        return tensor
+
+def _distributed_available():
+    return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
+    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
+    if '--local-rank' in cfg:
+        del cfg['--local-rank']
+    # added by Xavier -- delete '--local-rank' in multi-nodes training, don't know why there is such a keyword
+    scfg = OmegaConf.structured(fields(**cfg))
+    return scfg

models/modeling_magvitv2.py

@@ -0,0 +1,440 @@
+from dataclasses import dataclass, field
+import numpy as np
+import torch
+import torch.nn as nn
+from .common_modules import *
+from .modeling_utils import ConfigMixin, ModelMixin, register_to_config
+from .misc import *
+import math
+
+class Updateable:
+    def do_update_step(
+            self, epoch: int, global_step: int, on_load_weights: bool = False
+    ):
+        for attr in self.__dir__():
+            if attr.startswith("_"):
+                continue
+            try:
+                module = getattr(self, attr)
+            except:
+                continue  # ignore attributes like property, which can't be retrived using getattr?
+            if isinstance(module, Updateable):
+                module.do_update_step(
+                    epoch, global_step, on_load_weights=on_load_weights
+                )
+        self.update_step(epoch, global_step, on_load_weights=on_load_weights)
+
+    def do_update_step_end(self, epoch: int, global_step: int):
+        for attr in self.__dir__():
+            if attr.startswith("_"):
+                continue
+            try:
+                module = getattr(self, attr)
+            except:
+                continue  # ignore attributes like property, which can't be retrived using getattr?
+            if isinstance(module, Updateable):
+                module.do_update_step_end(epoch, global_step)
+        self.update_step_end(epoch, global_step)
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        # override this method to implement custom update logic
+        # if on_load_weights is True, you should be careful doing things related to model evaluations,
+        # as the models and tensors are not guarenteed to be on the same device
+        pass
+
+    def update_step_end(self, epoch: int, global_step: int):
+        pass
+
+class VQGANEncoder(ModelMixin, ConfigMixin):
+    @dataclass
+    class Config:
+        ch: int = 128
+        ch_mult: List[int] = field(default_factory=lambda: [1, 2, 2, 4, 4])
+        num_res_blocks: List[int] = field(default_factory=lambda: [4, 3, 4, 3, 4])
+        attn_resolutions: List[int] = field(default_factory=lambda: [5])
+        dropout: float = 0.0
+        in_ch: int = 3
+        out_ch: int = 3
+        resolution: int = 256
+        z_channels: int = 13
+        double_z: bool = False
+
+    def __init__(self,
+                 ch: int = 128,
+                 ch_mult: List[int] = [1, 2, 2, 4, 4],
+                 num_res_blocks: List[int] = [4, 3, 4, 3, 4],
+                 attn_resolutions: List[int] = [5],
+                 dropout: float = 0.0,
+                 in_ch: int = 3,
+                 out_ch: int = 3,
+                 resolution: int = 256,
+                 z_channels: int = 13,
+                 double_z: bool = False):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_ch = in_ch
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(
+            self.in_ch, self.ch, kernel_size=3, stride=1, padding=1
+        )
+
+        curr_res = self.resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = self.ch * in_ch_mult[i_level]
+            block_out = self.ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks[i_level]):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, True)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.quant_conv = torch.nn.Conv2d(z_channels, z_channels, 1)
+        # for param in self.parameters():
+        #     broadcast(param, src=0)
+
+    def forward(self, x):
+        # timestep embedding
+        temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks[i_level]):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        h = self.quant_conv(h)
+        return h
+
+
+class LFQuantizer(nn.Module):
+    def __init__(self, num_codebook_entry: int = -1,
+                 codebook_dim: int = 13,
+                 beta: float = 0.25,
+                 entropy_multiplier: float = 0.1,
+                 commit_loss_multiplier: float = 0.1, ):
+        super().__init__()
+        self.codebook_size = 2 ** codebook_dim
+        print(
+            f"Look-up free quantizer with codebook size: {self.codebook_size}"
+        )
+        self.e_dim = codebook_dim
+        self.beta = beta
+
+        indices = torch.arange(self.codebook_size)
+
+        binary = (
+                         indices.unsqueeze(1)
+                         >> torch.arange(codebook_dim - 1, -1, -1, dtype=torch.long)
+                 ) & 1
+
+        embedding = binary.float() * 2 - 1
+        self.register_buffer("embedding", embedding)
+        self.register_buffer(
+            "power_vals", 2 ** torch.arange(codebook_dim - 1, -1, -1)
+        )
+        self.commit_loss_multiplier = commit_loss_multiplier
+        self.entropy_multiplier = entropy_multiplier
+
+    def get_indices(self, z_q):
+        return (
+            (self.power_vals.reshape(1, -1, 1, 1) * (z_q > 0).float())
+            .sum(1, keepdim=True)
+            .long()
+        )
+
+    def get_codebook_entry(self, indices, shape=None):
+        if shape is None:
+            h, w = int(math.sqrt(indices.shape[-1])), int(math.sqrt(indices.shape[-1]))
+        else:
+            h, w = shape
+        b, _ = indices.shape
+        indices = indices.reshape(-1)
+        z_q = self.embedding[indices]
+        z_q = z_q.view(b, h, w, -1)
+
+        # reshape back to match original input shape
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+    def forward(self, z, get_code=False):
+        """
+        Inputs the output of the encoder network z and maps it to a discrete
+        one-hot vector that is the index of the closest embedding vector e_j
+        z (continuous) -> z_q (discrete)
+        z.shape = (batch, channel, height, width)
+        quantization pipeline:
+            1. get encoder input (B,C,H,W)
+            2. flatten input to (B*H*W,C)
+        """
+        if get_code:
+            return self.get_codebook_entry(z)
+
+        # reshape z -> (batch, height, width, channel) and flatten
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.e_dim)
+        ge_zero = (z_flattened > 0).float()
+        ones = torch.ones_like(z_flattened)
+        z_q = ones * ge_zero + -ones * (1 - ge_zero)
+
+        # preserve gradients
+        z_q = z_flattened + (z_q - z_flattened).detach()
+
+        # compute entropy loss
+        CatDist = torch.distributions.categorical.Categorical
+        logit = torch.stack(
+            [
+                -(z_flattened - torch.ones_like(z_q)).pow(2),
+                -(z_flattened - torch.ones_like(z_q) * -1).pow(2),
+            ],
+            dim=-1,
+        )
+        cat_dist = CatDist(logits=logit)
+        entropy = cat_dist.entropy().mean()
+        mean_prob = cat_dist.probs.mean(0)
+        mean_entropy = CatDist(probs=mean_prob).entropy().mean()
+
+        # compute loss for embedding
+        commit_loss = torch.mean(
+            (z_q.detach() - z_flattened) ** 2
+        ) + self.beta * torch.mean((z_q - z_flattened.detach()) ** 2)
+
+        # reshape back to match original input shape
+        z_q = z_q.view(z.shape)
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return {
+            "z": z_q,
+            "quantizer_loss": commit_loss * self.commit_loss_multiplier,
+            "entropy_loss": (entropy - mean_entropy) * self.entropy_multiplier,
+            "indices": self.get_indices(z_q),
+        }
+
+
+class VQGANDecoder(ModelMixin, ConfigMixin):
+    def __init__(self, ch: int = 128,
+                 ch_mult: List[int] = [1, 1, 2, 2, 4],
+                 num_res_blocks: List[int] = [4, 4, 3, 4, 3],
+                 attn_resolutions: List[int] = [5],
+                 dropout: float = 0.0,
+                 in_ch: int = 3,
+                 out_ch: int = 3,
+                 resolution: int = 256,
+                 z_channels: int = 13,
+                 double_z: bool = False):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_ch = in_ch
+        self.give_pre_end = False
+
+        self.z_channels = z_channels
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,) + tuple(ch_mult)
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = self.resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        print(
+            "Working with z of shape {} = {} dimensions.".format(
+                self.z_shape, np.prod(self.z_shape)
+            )
+        )
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(
+            z_channels, block_in, kernel_size=3, stride=1, padding=1
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks[i_level]):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, True)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(
+            block_in, out_ch, kernel_size=3, stride=1, padding=1
+        )
+        self.post_quant_conv = torch.nn.Conv2d(
+            z_channels, z_channels, 1
+        )
+
+
+    def forward(self, z):
+        # assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+        # timestep embedding
+        temb = None
+        output = dict()
+        z = self.post_quant_conv(z)
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks[i_level]):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        output["output"] = h
+        if self.give_pre_end:
+            return output
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        output["output"] = h
+        return output
+
+
+class MAGVITv2(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+            self,
+    ):
+        super().__init__()
+
+        self.encoder = VQGANEncoder()
+        self.decoder = VQGANDecoder()
+        self.quantize = LFQuantizer()
+
+    def forward(self, pixel_values, return_loss=False):
+        pass
+
+    def encode(self, pixel_values, return_loss=False):
+        hidden_states = self.encoder(pixel_values)
+        quantized_states = self.quantize(hidden_states)['z']
+        codebook_indices = self.quantize.get_indices(quantized_states).reshape(pixel_values.shape[0], -1)
+        output = (quantized_states, codebook_indices)
+        return output
+
+    def get_code(self, pixel_values):
+        hidden_states = self.encoder(pixel_values)
+        codebook_indices = self.quantize.get_indices(self.quantize(hidden_states)['z']).reshape(pixel_values.shape[0], -1)
+
+        return codebook_indices
+
+    def decode_code(self, codebook_indices, shape=None):
+        z_q = self.quantize.get_codebook_entry(codebook_indices, shape=shape)
+
+        reconstructed_pixel_values = self.decoder(z_q)["output"]
+        return reconstructed_pixel_values
+
+
+if __name__ == '__main__':
+    encoder = VQGANEncoder()
+    import ipdb
+    ipdb.set_trace()
+    print()

models/modeling_showo.py

@@ -0,0 +1,206 @@
+import torch
+import torch.nn.functional as F
+from transformers import AutoConfig
+from .modeling_utils import ConfigMixin, ModelMixin, register_to_config
+from .sampling import cosine_schedule, mask_by_random_topk
+from .phi import PhiForCausalLM
+
+try:
+    import xformers.ops as xops
+
+    is_xformers_available = True
+except ImportError:
+    is_xformers_available = False
+
+
+class Showo(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+            self,
+            w_clip_vit,
+            vocab_size,
+            llm_vocab_size,
+            llm_model_path='',
+            codebook_size=8192,
+            num_vq_tokens=256,
+            **kwargs,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.register_to_config(mask_token_id=vocab_size - 1)
+        config = AutoConfig.from_pretrained(llm_model_path)
+        self.showo = PhiForCausalLM(config)
+        self.showo.resize_token_embeddings(self.vocab_size)
+        self.output_size = self.vocab_size
+
+        if self.w_clip_vit:
+            self.mm_projector = torch.nn.Sequential(
+                torch.nn.Linear(1024, 2048),
+                torch.nn.GELU(),
+                torch.nn.Linear(2048, 2048)
+            )
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = True
+
+    def forward(
+            self,
+            input_ids,
+            input_embeddings=None,
+            attention_mask=None,
+            labels=None,
+            label_smoothing=0.0,
+            config=None,
+            labels_mask_text=None,
+            labels_mask_image=None,
+            **kwargs,
+    ):
+
+        if input_embeddings is None:
+            logits = self.showo(input_ids=input_ids, attention_mask=attention_mask)['logits']
+        else:
+            logits = self.showo(inputs_embeds=input_embeddings, attention_mask=attention_mask)['logits']
+
+        if labels is not None:
+            raise NotImplementedError
+
+        return logits
+
+    def t2i_generate(
+            self,
+            input_ids: torch.LongTensor = None,
+            uncond_input_ids: torch.LongTensor = None,
+            attention_mask=None,
+            temperature=1.0,
+            timesteps=18,  # ideal number of steps is 18 in maskgit paper
+            guidance_scale=0,
+            noise_schedule=cosine_schedule,
+            generator: torch.Generator = None,
+            uni_prompting=None,
+            config=None,
+            **kwargs,
+    ):
+        """
+        Generate 1:1 similar to the original MaskGit repo
+        https://github.com/google-research/maskgit/blob/main/maskgit/libml/parallel_decode.py#L79
+        """
+        # begin with all image token ids masked
+        mask_token_id = self.config.mask_token_id
+        seq_len = config.model.showo.num_vq_tokens
+
+        input_ids_minus_lm_vocab_size = input_ids[:, -(seq_len + 1):-1].clone()
+        input_ids_minus_lm_vocab_size = torch.where(input_ids_minus_lm_vocab_size == mask_token_id,
+                                                    mask_token_id,
+                                                    input_ids_minus_lm_vocab_size - config.model.showo.llm_vocab_size - 10)
+        # import ipdb
+        # ipdb.set_trace()
+        if uncond_input_ids is not None:
+            uncond_prefix = uncond_input_ids[:, :config.dataset.preprocessing.max_seq_length + 1]
+
+        for step in range(timesteps):
+            if uncond_input_ids is not None and guidance_scale > 0:
+                uncond_input_ids = torch.cat(
+                    [uncond_prefix, input_ids[:, config.dataset.preprocessing.max_seq_length + 1:]], dim=1)
+                model_input = torch.cat([input_ids, uncond_input_ids])
+                cond_logits, uncond_logits = self(model_input, attention_mask=attention_mask).chunk(2)
+                # logits = uncond_logits + guidance_scale * (cond_logits - uncond_logits)
+                # it seems that muse has different cfg setting
+                logits = (1 + guidance_scale) * cond_logits - guidance_scale * uncond_logits
+                logits = logits[:, -(seq_len + 1):-1, config.model.showo.llm_vocab_size + 10:-1]
+            else:
+                logits = self(input_ids, attention_mask=attention_mask)
+                logits = logits[:, -(seq_len + 1):-1, config.model.showo.llm_vocab_size + 10:-1]
+
+            probs = logits.softmax(dim=-1)
+            sampled = probs.reshape(-1, logits.size(-1))
+            sampled_ids = torch.multinomial(sampled, 1, generator=generator)[:, 0].view(*logits.shape[:-1])
+
+            unknown_map = input_ids_minus_lm_vocab_size == mask_token_id
+            sampled_ids = torch.where(unknown_map, sampled_ids, input_ids_minus_lm_vocab_size)
+            # Defines the mask ratio for the next round. The number to mask out is
+            # determined by mask_ratio * unknown_number_in_the_beginning.
+            ratio = 1.0 * (step + 1) / timesteps
+            mask_ratio = noise_schedule(torch.tensor(ratio))
+            # Computes the probabilities of each selected tokens.
+            selected_probs = torch.gather(probs, -1, sampled_ids.long()[..., None])
+            selected_probs = selected_probs.squeeze(-1)
+
+            # Ignores the tokens given in the input by overwriting their confidence.
+            selected_probs = torch.where(unknown_map, selected_probs, torch.finfo(selected_probs.dtype).max)
+            # Gets mask lens for each sample in the batch according to the mask ratio.
+            mask_len = (seq_len * mask_ratio).floor().unsqueeze(0).to(logits.device)
+            # Keeps at least one of prediction in this round and also masks out at least
+            # one and for the next iteration
+            mask_len = torch.max(
+                torch.tensor([1], device=logits.device), torch.min(unknown_map.sum(dim=-1, keepdim=True) - 1, mask_len)
+            )
+            # Adds noise for randomness
+            temperature = temperature * (1.0 - ratio)
+            masking = mask_by_random_topk(mask_len, selected_probs, temperature, generator=generator)
+            # Masks tokens with lower confidence.
+            input_ids[:, -(seq_len + 1):-1] = torch.where(masking, mask_token_id,
+                                                          sampled_ids + config.model.showo.llm_vocab_size + 10)
+            input_ids_minus_lm_vocab_size = torch.where(masking, mask_token_id, sampled_ids)
+
+        return sampled_ids
+
+    @torch.no_grad()
+    def mmu_generate(self, idx=None, input_embeddings=None, attention_mask=None, max_new_tokens=100, temperature=1.0, top_k=None, eot_token=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        try:
+            device = idx.device
+        except:
+            device = input_embeddings.device
+
+        result = []
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            # idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            # logits, _ = self(idx_cond)
+            logits = self(idx, input_embeddings=input_embeddings, attention_mask=attention_mask)
+
+            L = attention_mask.shape[-1]
+            attention_mask = attention_mask.squeeze()
+            attention_mask_a = torch.hstack(
+                [
+                    attention_mask,  # L, L
+                    torch.zeros((L, 1)).to(device) + torch.finfo(logits.dtype).min,
+                ]
+            )
+            attention_mask_b = torch.vstack(
+                [
+                    attention_mask_a,  # L, L+1
+                    torch.hstack([attention_mask[-1, :], torch.tensor([0]).to(device)]).unsqueeze(0),
+                ]
+            )
+            attention_mask = attention_mask_b
+
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            result.append(idx_next[0][0])
+            # append sampled index to the running sequence and continue
+            if self.config.w_clip_vit:
+                idx_next_embeddings = self.showo.model.embed_tokens(idx_next)
+                input_embeddings = torch.cat([input_embeddings, idx_next_embeddings], dim=1)
+            else:
+                idx = torch.cat((idx, idx_next), dim=1)
+
+            if eot_token is not None and idx_next.cpu() == eot_token:
+                break
+
+        return result

models/modeling_utils.py

@@ -0,0 +1,1207 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import itertools
+import json
+import os
+import re
+from collections import OrderedDict
+from functools import partial
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import safetensors
+import torch
+from huggingface_hub import create_repo, split_torch_state_dict_into_shards
+from huggingface_hub.utils import validate_hf_hub_args
+from torch import Tensor, nn
+
+from diffusers import __version__
+from diffusers.utils import (
+    FLAX_WEIGHTS_NAME,
+    SAFE_WEIGHTS_INDEX_NAME,
+    WEIGHTS_INDEX_NAME,
+    _add_variant,
+    _get_checkpoint_shard_files,
+    _get_model_file,
+    deprecate,
+    is_accelerate_available,
+    is_torch_version,
+    logging,
+)
+
+CONFIG_NAME = "config.json"
+WEIGHTS_NAME = "pytorch_model.bin"
+SAFETENSORS_WEIGHTS_NAME = "pytorch_model.safetensors"
+HUGGINGFACE_CO_RESOLVE_ENDPOINT = "https://huggingface.co"
+
+from diffusers.utils.hub_utils import (
+    PushToHubMixin,
+    load_or_create_model_card,
+    populate_model_card,
+)
+from diffusers.models.model_loading_utils import (
+    _determine_device_map,
+    _fetch_index_file,
+    _load_state_dict_into_model,
+    load_model_dict_into_meta,
+    load_state_dict,
+)
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+logger = logging.get_logger(__name__)
+
+_REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}")
+
+
+if is_torch_version(">=", "1.9.0"):
+    _LOW_CPU_MEM_USAGE_DEFAULT = True
+else:
+    _LOW_CPU_MEM_USAGE_DEFAULT = False
+
+
+if is_accelerate_available():
+    import accelerate
+
+
+def get_parameter_device(parameter: torch.nn.Module) -> torch.device:
+    try:
+        parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
+        return next(parameters_and_buffers).device
+    except StopIteration:
+        # For torch.nn.DataParallel compatibility in PyTorch 1.5
+
+        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+            return tuples
+
+        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+        first_tuple = next(gen)
+        return first_tuple[1].device
+
+
+def get_parameter_dtype(parameter: torch.nn.Module) -> torch.dtype:
+    try:
+        params = tuple(parameter.parameters())
+        if len(params) > 0:
+            return params[0].dtype
+
+        buffers = tuple(parameter.buffers())
+        if len(buffers) > 0:
+            return buffers[0].dtype
+
+    except StopIteration:
+        # For torch.nn.DataParallel compatibility in PyTorch 1.5
+
+        def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
+            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
+            return tuples
+
+        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
+        first_tuple = next(gen)
+        return first_tuple[1].dtype
+
+
+class ModelMixin(torch.nn.Module, PushToHubMixin):
+    r"""
+    Base class for all models.
+
+    [`ModelMixin`] takes care of storing the model configuration and provides methods for loading, downloading and
+    saving models.
+
+        - **config_name** ([`str`]) -- Filename to save a model to when calling [`~models.ModelMixin.save_pretrained`].
+    """
+
+    config_name = CONFIG_NAME
+    _automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
+    _supports_gradient_checkpointing = False
+    _keys_to_ignore_on_load_unexpected = None
+    _no_split_modules = None
+
+    def __init__(self):
+        super().__init__()
+
+    def __getattr__(self, name: str) -> Any:
+        """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
+        config attributes directly. See https://github.com/huggingface/diffusers/pull/3129 We need to overwrite
+        __getattr__ here in addition so that we don't trigger `torch.nn.Module`'s __getattr__':
+        https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
+        """
+
+        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
+        is_attribute = name in self.__dict__
+
+        if is_in_config and not is_attribute:
+            deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'unet.config.{name}'."
+            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)
+            return self._internal_dict[name]
+
+        # call PyTorch's https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
+        return super().__getattr__(name)
+
+    @property
+    def is_gradient_checkpointing(self) -> bool:
+        """
+        Whether gradient checkpointing is activated for this model or not.
+        """
+        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
+
+    def enable_gradient_checkpointing(self) -> None:
+        """
+        Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
+        *checkpoint activations* in other frameworks).
+        """
+        if not self._supports_gradient_checkpointing:
+            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
+        self.apply(partial(self._set_gradient_checkpointing, value=True))
+
+    def disable_gradient_checkpointing(self) -> None:
+        """
+        Deactivates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or
+        *checkpoint activations* in other frameworks).
+        """
+        if self._supports_gradient_checkpointing:
+            self.apply(partial(self._set_gradient_checkpointing, value=False))
+
+    def set_use_npu_flash_attention(self, valid: bool) -> None:
+        r"""
+        Set the switch for the npu flash attention.
+        """
+
+        def fn_recursive_set_npu_flash_attention(module: torch.nn.Module):
+            if hasattr(module, "set_use_npu_flash_attention"):
+                module.set_use_npu_flash_attention(valid)
+
+            for child in module.children():
+                fn_recursive_set_npu_flash_attention(child)
+
+        for module in self.children():
+            if isinstance(module, torch.nn.Module):
+                fn_recursive_set_npu_flash_attention(module)
+
+    def enable_npu_flash_attention(self) -> None:
+        r"""
+        Enable npu flash attention from torch_npu
+
+        """
+        self.set_use_npu_flash_attention(True)
+
+    def disable_npu_flash_attention(self) -> None:
+        r"""
+        disable npu flash attention from torch_npu
+
+        """
+        self.set_use_npu_flash_attention(False)
+
+    def set_use_memory_efficient_attention_xformers(
+        self, valid: bool, attention_op: Optional[Callable] = None
+    ) -> None:
+        # Recursively walk through all the children.
+        # Any children which exposes the set_use_memory_efficient_attention_xformers method
+        # gets the message
+        def fn_recursive_set_mem_eff(module: torch.nn.Module):
+            if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+                module.set_use_memory_efficient_attention_xformers(valid, attention_op)
+
+            for child in module.children():
+                fn_recursive_set_mem_eff(child)
+
+        for module in self.children():
+            if isinstance(module, torch.nn.Module):
+                fn_recursive_set_mem_eff(module)
+
+    def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None) -> None:
+        r"""
+        Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
+
+        When this option is enabled, you should observe lower GPU memory usage and a potential speed up during
+        inference. Speed up during training is not guaranteed.
+
+        <Tip warning={true}>
+
+        ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes
+        precedent.
+
+        </Tip>
+
+        Parameters:
+            attention_op (`Callable`, *optional*):
+                Override the default `None` operator for use as `op` argument to the
+                [`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention)
+                function of xFormers.
+
+        Examples:
+
+        ```py
+        >>> import torch
+        >>> from diffusers import UNet2DConditionModel
+        >>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
+
+        >>> model = UNet2DConditionModel.from_pretrained(
+        ...     "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
+        ... )
+        >>> model = model.to("cuda")
+        >>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
+        ```
+        """
+        self.set_use_memory_efficient_attention_xformers(True, attention_op)
+
+    def disable_xformers_memory_efficient_attention(self) -> None:
+        r"""
+        Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
+        """
+        self.set_use_memory_efficient_attention_xformers(False)
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        is_main_process: bool = True,
+        save_function: Optional[Callable] = None,
+        safe_serialization: bool = True,
+        variant: Optional[str] = None,
+        max_shard_size: Union[int, str] = "10GB",
+        push_to_hub: bool = False,
+        **kwargs,
+    ):
+        """
+        Save a model and its configuration file to a directory so that it can be reloaded using the
+        [`~models.ModelMixin.from_pretrained`] class method.
+
+        Arguments:
+            save_directory (`str` or `os.PathLike`):
+                Directory to save a model and its configuration file to. Will be created if it doesn't exist.
+            is_main_process (`bool`, *optional*, defaults to `True`):
+                Whether the process calling this is the main process or not. Useful during distributed training and you
+                need to call this function on all processes. In this case, set `is_main_process=True` only on the main
+                process to avoid race conditions.
+            save_function (`Callable`):
+                The function to use to save the state dictionary. Useful during distributed training when you need to
+                replace `torch.save` with another method. Can be configured with the environment variable
+                `DIFFUSERS_SAVE_MODE`.
+            safe_serialization (`bool`, *optional*, defaults to `True`):
+                Whether to save the model using `safetensors` or the traditional PyTorch way with `pickle`.
+            variant (`str`, *optional*):
+                If specified, weights are saved in the format `pytorch_model.<variant>.bin`.
+            max_shard_size (`int` or `str`, defaults to `"10GB"`):
+                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
+                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5GB"`).
+                If expressed as an integer, the unit is bytes. Note that this limit will be decreased after a certain
+                period of time (starting from Oct 2024) to allow users to upgrade to the latest version of `diffusers`.
+                This is to establish a common default size for this argument across different libraries in the Hugging
+                Face ecosystem (`transformers`, and `accelerate`, for example).
+            push_to_hub (`bool`, *optional*, defaults to `False`):
+                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
+                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
+                namespace).
+            kwargs (`Dict[str, Any]`, *optional*):
+                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
+        """
+        if os.path.isfile(save_directory):
+            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
+            return
+
+        weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
+        weights_name = _add_variant(weights_name, variant)
+        weight_name_split = weights_name.split(".")
+        if len(weight_name_split) in [2, 3]:
+            weights_name_pattern = weight_name_split[0] + "{suffix}." + ".".join(weight_name_split[1:])
+        else:
+            raise ValueError(f"Invalid {weights_name} provided.")
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        if push_to_hub:
+            commit_message = kwargs.pop("commit_message", None)
+            private = kwargs.pop("private", False)
+            create_pr = kwargs.pop("create_pr", False)
+            token = kwargs.pop("token", None)
+            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+            repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
+
+        # Only save the model itself if we are using distributed training
+        model_to_save = self
+
+        # Attach architecture to the config
+        # Save the config
+        if is_main_process:
+            model_to_save.save_config(save_directory)
+
+        # Save the model
+        state_dict = model_to_save.state_dict()
+
+        # Save the model
+        state_dict_split = split_torch_state_dict_into_shards(
+            state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern
+        )
+
+        # Clean the folder from a previous save
+        if is_main_process:
+            for filename in os.listdir(save_directory):
+                if filename in state_dict_split.filename_to_tensors.keys():
+                    continue
+                full_filename = os.path.join(save_directory, filename)
+                if not os.path.isfile(full_filename):
+                    continue
+                weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "")
+                weights_without_ext = weights_without_ext.replace("{suffix}", "")
+                filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "")
+                # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
+                if (
+                    filename.startswith(weights_without_ext)
+                    and _REGEX_SHARD.fullmatch(filename_without_ext) is not None
+                ):
+                    os.remove(full_filename)
+
+        for filename, tensors in state_dict_split.filename_to_tensors.items():
+            shard = {tensor: state_dict[tensor] for tensor in tensors}
+            filepath = os.path.join(save_directory, filename)
+            if safe_serialization:
+                # At some point we will need to deal better with save_function (used for TPU and other distributed
+                # joyfulness), but for now this enough.
+                safetensors.torch.save_file(shard, filepath, metadata={"format": "pt"})
+            else:
+                torch.save(shard, filepath)
+
+        if state_dict_split.is_sharded:
+            index = {
+                "metadata": state_dict_split.metadata,
+                "weight_map": state_dict_split.tensor_to_filename,
+            }
+            save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
+            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
+            # Save the index as well
+            with open(save_index_file, "w", encoding="utf-8") as f:
+                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+                f.write(content)
+            logger.info(
+                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
+                f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "
+                f"index located at {save_index_file}."
+            )
+        else:
+            path_to_weights = os.path.join(save_directory, weights_name)
+            logger.info(f"Model weights saved in {path_to_weights}")
+
+        if push_to_hub:
+            # Create a new empty model card and eventually tag it
+            model_card = load_or_create_model_card(repo_id, token=token)
+            model_card = populate_model_card(model_card)
+            model_card.save(Path(save_directory, "README.md").as_posix())
+
+            self._upload_folder(
+                save_directory,
+                repo_id,
+                token=token,
+                commit_message=commit_message,
+                create_pr=create_pr,
+            )
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        r"""
+        Instantiate a pretrained PyTorch model from a pretrained model configuration.
+
+        The model is set in evaluation mode - `model.eval()` - by default, and dropout modules are deactivated. To
+        train the model, set it back in training mode with `model.train()`.
+
+        Parameters:
+            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
+                Can be either:
+
+                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
+                      the Hub.
+                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
+                      with [`~ModelMixin.save_pretrained`].
+
+            cache_dir (`Union[str, os.PathLike]`, *optional*):
+                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
+                is not used.
+            torch_dtype (`str` or `torch.dtype`, *optional*):
+                Override the default `torch.dtype` and load the model with another dtype. If `"auto"` is passed, the
+                dtype is automatically derived from the model's weights.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            output_loading_info (`bool`, *optional*, defaults to `False`):
+                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
+            local_files_only(`bool`, *optional*, defaults to `False`):
+                Whether to only load local model weights and configuration files or not. If set to `True`, the model
+                won't be downloaded from the Hub.
+            token (`str` or *bool*, *optional*):
+                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
+                `diffusers-cli login` (stored in `~/.huggingface`) is used.
+            revision (`str`, *optional*, defaults to `"main"`):
+                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
+                allowed by Git.
+            from_flax (`bool`, *optional*, defaults to `False`):
+                Load the model weights from a Flax checkpoint save file.
+            subfolder (`str`, *optional*, defaults to `""`):
+                The subfolder location of a model file within a larger model repository on the Hub or locally.
+            mirror (`str`, *optional*):
+                Mirror source to resolve accessibility issues if you're downloading a model in China. We do not
+                guarantee the timeliness or safety of the source, and you should refer to the mirror site for more
+                information.
+            device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
+                A map that specifies where each submodule should go. It doesn't need to be defined for each
+                parameter/buffer name; once a given module name is inside, every submodule of it will be sent to the
+                same device. Defaults to `None`, meaning that the model will be loaded on CPU.
+
+                Set `device_map="auto"` to have 🤗 Accelerate automatically compute the most optimized `device_map`. For
+                more information about each option see [designing a device
+                map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
+            max_memory (`Dict`, *optional*):
+                A dictionary device identifier for the maximum memory. Will default to the maximum memory available for
+                each GPU and the available CPU RAM if unset.
+            offload_folder (`str` or `os.PathLike`, *optional*):
+                The path to offload weights if `device_map` contains the value `"disk"`.
+            offload_state_dict (`bool`, *optional*):
+                If `True`, temporarily offloads the CPU state dict to the hard drive to avoid running out of CPU RAM if
+                the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to `True`
+                when there is some disk offload.
+            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
+                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
+                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
+                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
+                argument to `True` will raise an error.
+            variant (`str`, *optional*):
+                Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when
+                loading `from_flax`.
+            use_safetensors (`bool`, *optional*, defaults to `None`):
+                If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the
+                `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors`
+                weights. If set to `False`, `safetensors` weights are not loaded.
+
+        <Tip>
+
+        To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in with
+        `huggingface-cli login`. You can also activate the special
+        ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a
+        firewalled environment.
+
+        </Tip>
+
+        Example:
+
+        ```py
+        from diffusers import UNet2DConditionModel
+
+        unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
+        ```
+
+        If you get the error message below, you need to finetune the weights for your downstream task:
+
+        ```bash
+        Some weights of UNet2DConditionModel were not initialized from the model checkpoint at runwayml/stable-diffusion-v1-5 and are newly initialized because the shapes did not match:
+        - conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated
+        You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
+        ```
+        """
+        cache_dir = kwargs.pop("cache_dir", None)
+        ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
+        force_download = kwargs.pop("force_download", False)
+        from_flax = kwargs.pop("from_flax", False)
+        proxies = kwargs.pop("proxies", None)
+        output_loading_info = kwargs.pop("output_loading_info", False)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        subfolder = kwargs.pop("subfolder", None)
+        device_map = kwargs.pop("device_map", None)
+        max_memory = kwargs.pop("max_memory", None)
+        offload_folder = kwargs.pop("offload_folder", None)
+        offload_state_dict = kwargs.pop("offload_state_dict", False)
+        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
+        variant = kwargs.pop("variant", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        if low_cpu_mem_usage and not is_accelerate_available():
+            low_cpu_mem_usage = False
+            logger.warning(
+                "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
+                " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
+                " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
+                " install accelerate\n```\n."
+            )
+
+        if device_map is not None and not is_accelerate_available():
+            raise NotImplementedError(
+                "Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set"
+                " `device_map=None`. You can install accelerate with `pip install accelerate`."
+            )
+
+        # Check if we can handle device_map and dispatching the weights
+        if device_map is not None and not is_torch_version(">=", "1.9.0"):
+            raise NotImplementedError(
+                "Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
+                " `device_map=None`."
+            )
+
+        if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
+            raise NotImplementedError(
+                "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
+                " `low_cpu_mem_usage=False`."
+            )
+
+        if low_cpu_mem_usage is False and device_map is not None:
+            raise ValueError(
+                f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and"
+                " dispatching. Please make sure to set `low_cpu_mem_usage=True`."
+            )
+
+        # change device_map into a map if we passed an int, a str or a torch.device
+        if isinstance(device_map, torch.device):
+            device_map = {"": device_map}
+        elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
+            try:
+                device_map = {"": torch.device(device_map)}
+            except RuntimeError:
+                raise ValueError(
+                    "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or "
+                    f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}."
+                )
+        elif isinstance(device_map, int):
+            if device_map < 0:
+                raise ValueError(
+                    "You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' "
+                )
+            else:
+                device_map = {"": device_map}
+
+        if device_map is not None:
+            if low_cpu_mem_usage is None:
+                low_cpu_mem_usage = True
+            elif not low_cpu_mem_usage:
+                raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`")
+
+        if low_cpu_mem_usage:
+            if device_map is not None and not is_torch_version(">=", "1.10"):
+                # The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info.
+                raise ValueError("`low_cpu_mem_usage` and `device_map` require PyTorch >= 1.10.")
+
+        # Load config if we don't provide a configuration
+        config_path = pretrained_model_name_or_path
+
+        user_agent = {
+            "diffusers": __version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # load config
+        config, unused_kwargs, commit_hash = cls.load_config(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            return_commit_hash=True,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            **kwargs,
+        )
+
+        # Determine if we're loading from a directory of sharded checkpoints.
+        is_sharded = False
+        index_file = None
+        is_local = os.path.isdir(pretrained_model_name_or_path)
+        index_file = _fetch_index_file(
+            is_local=is_local,
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            subfolder=subfolder or "",
+            use_safetensors=use_safetensors,
+            cache_dir=cache_dir,
+            variant=variant,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            user_agent=user_agent,
+            commit_hash=commit_hash,
+        )
+        if index_file is not None and index_file.is_file():
+            is_sharded = True
+
+        if is_sharded and from_flax:
+            raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.")
+
+        # load model
+        model_file = None
+        if from_flax:
+            model_file = _get_model_file(
+                pretrained_model_name_or_path,
+                weights_name=FLAX_WEIGHTS_NAME,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder=subfolder,
+                user_agent=user_agent,
+                commit_hash=commit_hash,
+            )
+            model = cls.from_config(config, **unused_kwargs)
+
+            # Convert the weights
+            from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model
+
+            model = load_flax_checkpoint_in_pytorch_model(model, model_file)
+        else:
+            if is_sharded:
+                sharded_ckpt_cached_folder, sharded_metadata = _get_checkpoint_shard_files(
+                    pretrained_model_name_or_path,
+                    index_file,
+                    cache_dir=cache_dir,
+                    proxies=proxies,
+                    local_files_only=local_files_only,
+                    token=token,
+                    user_agent=user_agent,
+                    revision=revision,
+                    subfolder=subfolder or "",
+                )
+
+            elif use_safetensors and not is_sharded:
+                try:
+                    model_file = _get_model_file(
+                        pretrained_model_name_or_path,
+                        weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
+                        cache_dir=cache_dir,
+                        force_download=force_download,
+                        proxies=proxies,
+                        local_files_only=local_files_only,
+                        token=token,
+                        revision=revision,
+                        subfolder=subfolder,
+                        user_agent=user_agent,
+                        commit_hash=commit_hash,
+                    )
+
+                except IOError as e:
+                    logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
+                    if not allow_pickle:
+                        raise
+                    logger.warning(
+                        "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
+                    )
+
+            if model_file is None and not is_sharded:
+                model_file = _get_model_file(
+                    pretrained_model_name_or_path,
+                    weights_name=_add_variant(WEIGHTS_NAME, variant),
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                    subfolder=subfolder,
+                    user_agent=user_agent,
+                    commit_hash=commit_hash,
+                )
+
+            if low_cpu_mem_usage:
+                # Instantiate model with empty weights
+                with accelerate.init_empty_weights():
+                    model = cls.from_config(config, **unused_kwargs)
+
+                # if device_map is None, load the state dict and move the params from meta device to the cpu
+                if device_map is None and not is_sharded:
+                    param_device = "cpu"
+                    state_dict = load_state_dict(model_file, variant=variant)
+                    model._convert_deprecated_attention_blocks(state_dict)
+                    # move the params from meta device to cpu
+                    missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+                    if len(missing_keys) > 0:
+                        raise ValueError(
+                            f"Cannot load {cls} from {pretrained_model_name_or_path} because the following keys are"
+                            f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+                            " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+                            " those weights or else make sure your checkpoint file is correct."
+                        )
+
+                    unexpected_keys = load_model_dict_into_meta(
+                        model,
+                        state_dict,
+                        device=param_device,
+                        dtype=torch_dtype,
+                        model_name_or_path=pretrained_model_name_or_path,
+                    )
+
+                    if cls._keys_to_ignore_on_load_unexpected is not None:
+                        for pat in cls._keys_to_ignore_on_load_unexpected:
+                            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+                    if len(unexpected_keys) > 0:
+                        logger.warning(
+                            f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+                        )
+
+                else:  # else let accelerate handle loading and dispatching.
+                    # Load weights and dispatch according to the device_map
+                    # by default the device_map is None and the weights are loaded on the CPU
+                    force_hook = True
+                    device_map = _determine_device_map(model, device_map, max_memory, torch_dtype)
+                    if device_map is None and is_sharded:
+                        # we load the parameters on the cpu
+                        device_map = {"": "cpu"}
+                        force_hook = False
+                    try:
+                        accelerate.load_checkpoint_and_dispatch(
+                            model,
+                            model_file if not is_sharded else index_file,
+                            device_map,
+                            max_memory=max_memory,
+                            offload_folder=offload_folder,
+                            offload_state_dict=offload_state_dict,
+                            dtype=torch_dtype,
+                            force_hooks=force_hook,
+                            strict=True,
+                        )
+                    except AttributeError as e:
+                        # When using accelerate loading, we do not have the ability to load the state
+                        # dict and rename the weight names manually. Additionally, accelerate skips
+                        # torch loading conventions and directly writes into `module.{_buffers, _parameters}`
+                        # (which look like they should be private variables?), so we can't use the standard hooks
+                        # to rename parameters on load. We need to mimic the original weight names so the correct
+                        # attributes are available. After we have loaded the weights, we convert the deprecated
+                        # names to the new non-deprecated names. Then we _greatly encourage_ the user to convert
+                        # the weights so we don't have to do this again.
+
+                        if "'Attention' object has no attribute" in str(e):
+                            logger.warning(
+                                f"Taking `{str(e)}` while using `accelerate.load_checkpoint_and_dispatch` to mean {pretrained_model_name_or_path}"
+                                " was saved with deprecated attention block weight names. We will load it with the deprecated attention block"
+                                " names and convert them on the fly to the new attention block format. Please re-save the model after this conversion,"
+                                " so we don't have to do the on the fly renaming in the future. If the model is from a hub checkpoint,"
+                                " please also re-upload it or open a PR on the original repository."
+                            )
+                            model._temp_convert_self_to_deprecated_attention_blocks()
+                            accelerate.load_checkpoint_and_dispatch(
+                                model,
+                                model_file if not is_sharded else index_file,
+                                device_map,
+                                max_memory=max_memory,
+                                offload_folder=offload_folder,
+                                offload_state_dict=offload_state_dict,
+                                dtype=torch_dtype,
+                                force_hooks=force_hook,
+                                strict=True,
+                            )
+                            model._undo_temp_convert_self_to_deprecated_attention_blocks()
+                        else:
+                            raise e
+
+                loading_info = {
+                    "missing_keys": [],
+                    "unexpected_keys": [],
+                    "mismatched_keys": [],
+                    "error_msgs": [],
+                }
+            else:
+                model = cls.from_config(config, **unused_kwargs)
+
+                state_dict = load_state_dict(model_file, variant=variant)
+                model._convert_deprecated_attention_blocks(state_dict)
+
+                model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
+                    model,
+                    state_dict,
+                    model_file,
+                    pretrained_model_name_or_path,
+                    ignore_mismatched_sizes=ignore_mismatched_sizes,
+                )
+
+                loading_info = {
+                    "missing_keys": missing_keys,
+                    "unexpected_keys": unexpected_keys,
+                    "mismatched_keys": mismatched_keys,
+                    "error_msgs": error_msgs,
+                }
+
+        if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
+            raise ValueError(
+                f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
+            )
+        elif torch_dtype is not None:
+            model = model.to(torch_dtype)
+
+        model.register_to_config(_name_or_path=pretrained_model_name_or_path)
+
+        # Set model in evaluation mode to deactivate DropOut modules by default
+        model.eval()
+        if output_loading_info:
+            return model, loading_info
+
+        return model
+
+    @classmethod
+    def _load_pretrained_model(
+        cls,
+        model,
+        state_dict: OrderedDict,
+        resolved_archive_file,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        ignore_mismatched_sizes: bool = False,
+    ):
+        # Retrieve missing & unexpected_keys
+        model_state_dict = model.state_dict()
+        loaded_keys = list(state_dict.keys())
+
+        expected_keys = list(model_state_dict.keys())
+
+        original_loaded_keys = loaded_keys
+
+        missing_keys = list(set(expected_keys) - set(loaded_keys))
+        unexpected_keys = list(set(loaded_keys) - set(expected_keys))
+
+        # Make sure we are able to load base models as well as derived models (with heads)
+        model_to_load = model
+
+        def _find_mismatched_keys(
+            state_dict,
+            model_state_dict,
+            loaded_keys,
+            ignore_mismatched_sizes,
+        ):
+            mismatched_keys = []
+            if ignore_mismatched_sizes:
+                for checkpoint_key in loaded_keys:
+                    model_key = checkpoint_key
+
+                    if (
+                        model_key in model_state_dict
+                        and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
+                    ):
+                        mismatched_keys.append(
+                            (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
+                        )
+                        del state_dict[checkpoint_key]
+            return mismatched_keys
+
+        if state_dict is not None:
+            # Whole checkpoint
+            mismatched_keys = _find_mismatched_keys(
+                state_dict,
+                model_state_dict,
+                original_loaded_keys,
+                ignore_mismatched_sizes,
+            )
+            error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
+
+        if len(error_msgs) > 0:
+            error_msg = "\n\t".join(error_msgs)
+            if "size mismatch" in error_msg:
+                error_msg += (
+                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
+                )
+            raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
+
+        if len(unexpected_keys) > 0:
+            logger.warning(
+                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
+                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
+                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task"
+                " or with another architecture (e.g. initializing a BertForSequenceClassification model from a"
+                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
+                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly"
+                " identical (initializing a BertForSequenceClassification model from a"
+                " BertForSequenceClassification model)."
+            )
+        else:
+            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
+        if len(missing_keys) > 0:
+            logger.warning(
+                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
+                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
+                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
+            )
+        elif len(mismatched_keys) == 0:
+            logger.info(
+                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
+                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the"
+                f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions"
+                " without further training."
+            )
+        if len(mismatched_keys) > 0:
+            mismatched_warning = "\n".join(
+                [
+                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
+                    for key, shape1, shape2 in mismatched_keys
+                ]
+            )
+            logger.warning(
+                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
+                f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
+                f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be"
+                " able to use it for predictions and inference."
+            )
+
+        return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
+
+    @classmethod
+    def _get_signature_keys(cls, obj):
+        parameters = inspect.signature(obj.__init__).parameters
+        required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
+        optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
+        expected_modules = set(required_parameters.keys()) - {"self"}
+
+        return expected_modules, optional_parameters
+
+    # Adapted from `transformers` modeling_utils.py
+    def _get_no_split_modules(self, device_map: str):
+        """
+        Get the modules of the model that should not be spit when using device_map. We iterate through the modules to
+        get the underlying `_no_split_modules`.
+
+        Args:
+            device_map (`str`):
+                The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"]
+
+        Returns:
+            `List[str]`: List of modules that should not be split
+        """
+        _no_split_modules = set()
+        modules_to_check = [self]
+        while len(modules_to_check) > 0:
+            module = modules_to_check.pop(-1)
+            # if the module does not appear in _no_split_modules, we also check the children
+            if module.__class__.__name__ not in _no_split_modules:
+                if isinstance(module, ModelMixin):
+                    if module._no_split_modules is None:
+                        raise ValueError(
+                            f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model "
+                            "class needs to implement the `_no_split_modules` attribute."
+                        )
+                    else:
+                        _no_split_modules = _no_split_modules | set(module._no_split_modules)
+                modules_to_check += list(module.children())
+        return list(_no_split_modules)
+
+    @property
+    def device(self) -> torch.device:
+        """
+        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+        device).
+        """
+        return get_parameter_device(self)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
+        """
+        return get_parameter_dtype(self)
+
+    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
+        """
+        Get number of (trainable or non-embedding) parameters in the module.
+
+        Args:
+            only_trainable (`bool`, *optional*, defaults to `False`):
+                Whether or not to return only the number of trainable parameters.
+            exclude_embeddings (`bool`, *optional*, defaults to `False`):
+                Whether or not to return only the number of non-embedding parameters.
+
+        Returns:
+            `int`: The number of parameters.
+
+        Example:
+
+        ```py
+        from diffusers import UNet2DConditionModel
+
+        model_id = "runwayml/stable-diffusion-v1-5"
+        unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet")
+        unet.num_parameters(only_trainable=True)
+        859520964
+        ```
+        """
+
+        if exclude_embeddings:
+            embedding_param_names = [
+                f"{name}.weight"
+                for name, module_type in self.named_modules()
+                if isinstance(module_type, torch.nn.Embedding)
+            ]
+            non_embedding_parameters = [
+                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
+            ]
+            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
+        else:
+            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
+
+    def _convert_deprecated_attention_blocks(self, state_dict: OrderedDict) -> None:
+        deprecated_attention_block_paths = []
+
+        def recursive_find_attn_block(name, module):
+            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
+                deprecated_attention_block_paths.append(name)
+
+            for sub_name, sub_module in module.named_children():
+                sub_name = sub_name if name == "" else f"{name}.{sub_name}"
+                recursive_find_attn_block(sub_name, sub_module)
+
+        recursive_find_attn_block("", self)
+
+        # NOTE: we have to check if the deprecated parameters are in the state dict
+        # because it is possible we are loading from a state dict that was already
+        # converted
+
+        for path in deprecated_attention_block_paths:
+            # group_norm path stays the same
+
+            # query -> to_q
+            if f"{path}.query.weight" in state_dict:
+                state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")
+            if f"{path}.query.bias" in state_dict:
+                state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")
+
+            # key -> to_k
+            if f"{path}.key.weight" in state_dict:
+                state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")
+            if f"{path}.key.bias" in state_dict:
+                state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")
+
+            # value -> to_v
+            if f"{path}.value.weight" in state_dict:
+                state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")
+            if f"{path}.value.bias" in state_dict:
+                state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")
+
+            # proj_attn -> to_out.0
+            if f"{path}.proj_attn.weight" in state_dict:
+                state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")
+            if f"{path}.proj_attn.bias" in state_dict:
+                state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")
+
+    def _temp_convert_self_to_deprecated_attention_blocks(self) -> None:
+        deprecated_attention_block_modules = []
+
+        def recursive_find_attn_block(module):
+            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
+                deprecated_attention_block_modules.append(module)
+
+            for sub_module in module.children():
+                recursive_find_attn_block(sub_module)
+
+        recursive_find_attn_block(self)
+
+        for module in deprecated_attention_block_modules:
+            module.query = module.to_q
+            module.key = module.to_k
+            module.value = module.to_v
+            module.proj_attn = module.to_out[0]
+
+            # We don't _have_ to delete the old attributes, but it's helpful to ensure
+            # that _all_ the weights are loaded into the new attributes and we're not
+            # making an incorrect assumption that this model should be converted when
+            # it really shouldn't be.
+            del module.to_q
+            del module.to_k
+            del module.to_v
+            del module.to_out
+
+    def _undo_temp_convert_self_to_deprecated_attention_blocks(self) -> None:
+        deprecated_attention_block_modules = []
+
+        def recursive_find_attn_block(module) -> None:
+            if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:
+                deprecated_attention_block_modules.append(module)
+
+            for sub_module in module.children():
+                recursive_find_attn_block(sub_module)
+
+        recursive_find_attn_block(self)
+
+        for module in deprecated_attention_block_modules:
+            module.to_q = module.query
+            module.to_k = module.key
+            module.to_v = module.value
+            module.to_out = nn.ModuleList([module.proj_attn, nn.Dropout(module.dropout)])
+
+            del module.query
+            del module.key
+            del module.value
+            del module.proj_attn
+
+
+class LegacyModelMixin(ModelMixin):
+    r"""
+    A subclass of `ModelMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
+    pipeline-specific classes (like `DiTTransformer2DModel`).
+    """
+
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        # To prevent dependency import problem.
+        from diffusers.models.model_loading_utils import _fetch_remapped_cls_from_config
+
+        # Create a copy of the kwargs so that we don't mess with the keyword arguments in the downstream calls.
+        kwargs_copy = kwargs.copy()
+
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", None)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+
+        # Load config if we don't provide a configuration
+        config_path = pretrained_model_name_or_path
+
+        user_agent = {
+            "diffusers": __version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # load config
+        config, _, _ = cls.load_config(
+            config_path,
+            cache_dir=cache_dir,
+            return_unused_kwargs=True,
+            return_commit_hash=True,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            subfolder=subfolder,
+            user_agent=user_agent,
+            **kwargs,
+        )
+        # resolve remapping
+        remapped_class = _fetch_remapped_cls_from_config(config, cls)
+
+        return remapped_class.from_pretrained(pretrained_model_name_or_path, **kwargs_copy)

models/phi.py

@@ -0,0 +1,1489 @@
+# coding=utf-8
+# Copyright 2023 Microsoft and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PyTorch Phi model."""
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from packaging import version
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    get_torch_version,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.models.phi.configuration_phi import PhiConfig
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "microsoft/phi-1"
+_CONFIG_FOR_DOC = "PhiConfig"
+
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Phi
+class PhiRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
+        )
+
+
+# Copied from transformers.models.falcon.modeling_falcon.FalconLinearScalingRotaryEmbedding with Falcon->Phi
+class PhiLinearScalingRotaryEmbedding(PhiRotaryEmbedding):
+    """PhiRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+        t = t / self.scaling_factor
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+# Copied from transformers.models.falcon.modeling_falcon.FalconDynamicNTKScalingRotaryEmbedding with Falcon->Phi
+class PhiDynamicNTKScalingRotaryEmbedding(PhiRotaryEmbedding):
+    """PhiRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+
+        freqs = torch.outer(t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.mixtral.modeling_mixtral.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+            used to pass offsetted position ids when working with a KV-cache.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Phi
+class PhiMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv with llama->phi
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class PhiAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: PhiConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.partial_rotary_factor = config.partial_rotary_factor
+        self.is_causal = True
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
+        self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=True)
+
+        self.qk_layernorm = config.qk_layernorm
+        if self.qk_layernorm:
+            self.q_layernorm = nn.LayerNorm(
+                config.hidden_size // self.num_heads, eps=config.layer_norm_eps, elementwise_affine=True
+            )
+            self.k_layernorm = nn.LayerNorm(
+                config.hidden_size // self.num_heads, eps=config.layer_norm_eps, elementwise_affine=True
+            )
+
+        self._init_rope()
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = PhiRotaryEmbedding(
+                int(self.partial_rotary_factor * self.head_dim),
+                max_position_embeddings=self.max_position_embeddings,
+                base=self.rope_theta,
+            )
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            if scaling_type == "linear":
+                self.rotary_emb = PhiLinearScalingRotaryEmbedding(
+                    int(self.partial_rotary_factor * self.head_dim),
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            elif scaling_type == "dynamic":
+                self.rotary_emb = PhiDynamicNTKScalingRotaryEmbedding(
+                    int(self.partial_rotary_factor * self.head_dim),
+                    max_position_embeddings=self.max_position_embeddings,
+                    scaling_factor=scaling_factor,
+                    base=self.rope_theta,
+                )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_layernorm:
+            query_states = self.q_layernorm(query_states)
+            key_states = self.k_layernorm(key_states)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        # Partial rotary embedding
+        query_rot, query_pass = (
+            query_states[..., : self.rotary_emb.dim],
+            query_states[..., self.rotary_emb.dim :],
+        )
+        key_rot, key_pass = (
+            key_states[..., : self.rotary_emb.dim],
+            key_states[..., self.rotary_emb.dim :],
+        )
+        # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
+        query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
+
+        # [batch_size, seq_length, num_heads, head_dim]
+        query_states = torch.cat((query_rot, query_pass), dim=-1)
+        key_states = torch.cat((key_rot, key_pass), dim=-1)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
+        attn_weights = torch.matmul(
+            query_states.to(torch.float32), key_states.to(torch.float32).transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.dense(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class PhiFlashAttention2(PhiAttention):
+    """
+    Phi flash attention module. This module inherits from `PhiAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # PhiFlashAttention2 attention does not support output_attentions
+
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_layernorm:
+            query_states = self.q_layernorm(query_states)
+            key_states = self.k_layernorm(key_states)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        # Partial rotary embedding
+        query_rot, query_pass = (
+            query_states[..., : self.rotary_emb.dim],
+            query_states[..., self.rotary_emb.dim :],
+        )
+        key_rot, key_pass = (
+            key_states[..., : self.rotary_emb.dim],
+            key_states[..., self.rotary_emb.dim :],
+        )
+        # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
+        query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
+
+        # [batch_size, seq_length, num_heads, head_dim]
+        query_states = torch.cat((query_rot, query_pass), dim=-1)
+        key_states = torch.cat((key_rot, key_pass), dim=-1)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+        # to be able to avoid many of these transpose/reshape/view.
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        attn_dropout = self.attention_dropout if self.training else 0.0
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in the correct dtype just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32.
+
+        if query_states.dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=attn_dropout, softmax_scale=None
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.dense(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
+    def _flash_attention_forward(
+        self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`float`):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        if not self._flash_attn_uses_top_left_mask:
+            causal = self.is_causal
+        else:
+            # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
+            causal = self.is_causal and query_length != 1
+
+        # Contains at least one padding token in the sequence
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
+    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+class PhiSdpaAttention(PhiAttention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.require_contiguous_qkv = version.parse(get_torch_version()) < version.parse("2.2.0")
+
+    """
+    SDPA attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `PhiAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from PhiAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "PhiModel is using PhiSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not "
+                "support `output_attentions=True`. Falling back to the manual attention implementation, but specifying "
+                "the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can "
+                'be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_layernorm:
+            query_states = self.q_layernorm(query_states)
+            key_states = self.k_layernorm(key_states)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        # Partial rotary embedding
+        query_rot, query_pass = (
+            query_states[..., : self.rotary_emb.dim],
+            query_states[..., self.rotary_emb.dim :],
+        )
+        key_rot, key_pass = (
+            key_states[..., : self.rotary_emb.dim],
+            key_states[..., self.rotary_emb.dim :],
+        )
+        # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
+        query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids)
+
+        # [batch_size, seq_length, num_heads, head_dim]
+        query_states = torch.cat((query_rot, query_pass), dim=-1)
+        key_states = torch.cat((key_rot, key_pass), dim=-1)
+
+        if past_key_value is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "partial_rotation_size": self.rotary_emb.dim}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom
+        # attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0.
+        # Reference: https://github.com/pytorch/pytorch/issues/112577
+        if self.require_contiguous_qkv and query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+        is_causal = True if self.is_causal and attention_mask is None and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.dense(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+PHI_ATTENTION_CLASSES = {
+    "eager": PhiAttention,
+    "flash_attention_2": PhiFlashAttention2,
+    "sdpa": PhiSdpaAttention,
+}
+
+
+class PhiDecoderLayer(nn.Module):
+    def __init__(self, config: PhiConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = PHI_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
+        self.mlp = PhiMLP(config)
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`):
+                input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
+                Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
+                `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        attn_outputs, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        attn_outputs = self.resid_dropout(attn_outputs)
+
+        feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))
+        hidden_states = attn_outputs + feed_forward_hidden_states + residual
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+PHI_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`PhiConfig`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+    "The bare Phi Model outputting raw hidden-states without any specific head on top.",
+    PHI_START_DOCSTRING,
+)
+class PhiPreTrainedModel(PreTrainedModel):
+    config_class = PhiConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["PhiDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+PHI_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Phi Model outputting raw hidden-states without any specific head on top.",
+    PHI_START_DOCSTRING,
+)
+class PhiModel(PhiPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`PhiDecoderLayer`]
+
+    Args:
+        config: PhiConfig
+    """
+
+    def __init__(self, config: PhiConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.embed_dropout = nn.Dropout(config.embd_pdrop)
+        print("attention implementation: ", config._attn_implementation)
+        self.layers = nn.ModuleList(
+            [PhiDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self._use_sdpa = config._attn_implementation == "sdpa"
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(PHI_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        past_key_values_length = 0
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        if use_cache:
+            use_legacy_cache = not isinstance(past_key_values, Cache)
+            if use_legacy_cache:
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+            past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        inputs_embeds = self.embed_dropout(inputs_embeds)
+        # commented by Xavier
+        # Attention mask.
+        # if self._use_flash_attention_2:
+        #     # 2d mask is passed through the layers
+        #     attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        # elif self._use_sdpa and not output_attentions:
+        #     attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+        #         attention_mask,
+        #         (batch_size, seq_length),
+        #         inputs_embeds,
+        #         past_key_values_length,
+        #     )
+        # else:
+        #     # 4d mask is passed through the layers
+        #     attention_mask = _prepare_4d_causal_attention_mask(
+        #         attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+        #     )
+        # commented by Xavier
+
+        hidden_states = inputs_embeds
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.final_layernorm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class PhiForCausalLM(PhiPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    def __init__(self, config):
+        super().__init__(config)
+        config.qk_layernorm = True
+        config.use_cache = False
+        self.model = PhiModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(PHI_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, PhiForCausalLM
+
+        >>> model = PhiForCausalLM.from_pretrained("microsoft/phi-1")
+        >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1")
+
+        >>> prompt = "This is an example script ."
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        'This is an example script .\n\n\n\nfrom typing import List\n\ndef find_most_common_letter(words: List[str'
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonForCausalLM.prepare_inputs_for_generation
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = past_key_values.get_max_length()
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM._reorder_cache
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+
+@add_start_docstrings(
+    """
+    The PhiModel with a sequence classification head on top (linear layer).
+
+    [`PhiForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    PHI_START_DOCSTRING,
+)
+# Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->PHI,Llama->Phi with self.transformer->self.model, transformer_outputs->model_outputs
+class PhiForSequenceClassification(PhiPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = PhiModel(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    @add_start_docstrings_to_model_forward(PHI_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        model_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = model_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + model_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=model_outputs.past_key_values,
+            hidden_states=model_outputs.hidden_states,
+            attentions=model_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    PhiModel with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    PHI_START_DOCSTRING,
+)
+# Copied from transformers.models.mpt.modeling_mpt.MptForTokenClassification with MPT->PHI,Mpt->Phi,self.transformer->self.model,transformer_outputs->model_outputs
+class PhiForTokenClassification(PhiPreTrainedModel):
+    def __init__(self, config: PhiConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.model = PhiModel(config)
+        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+            classifier_dropout = config.classifier_dropout
+        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(PHI_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        model_outputs = self.model(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = model_outputs[0]
+        hidden_states = self.dropout(hidden_states)
+        logits = self.classifier(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            batch_size, seq_length = labels.shape
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
+            )
+
+        if not return_dict:
+            output = (logits,) + model_outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=model_outputs.hidden_states,
+            attentions=model_outputs.attentions,
+        )

models/sampling.py

@@ -0,0 +1,118 @@
+# Adapted from https://github.com/lucidrains/muse-maskgit-pytorch
+
+import math
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+
+def log(t, eps=1e-20):
+    return torch.log(t.clamp(min=eps))
+
+
+def gumbel_noise(t, generator=None):
+    noise = torch.zeros_like(t).uniform_(0, 1, generator=generator)
+    return -log(-log(noise))
+
+
+def gumbel_sample(t, temperature=1.0, dim=-1, generator=None):
+    return ((t / max(temperature, 1e-10)) + gumbel_noise(t, generator=generator)).argmax(dim=dim)
+
+
+def top_k(logits, thres=0.9):
+    k = math.ceil((1 - thres) * logits.shape[-1])
+    val, ind = logits.topk(k, dim=-1)
+    probs = torch.full_like(logits, float("-inf"))
+    probs.scatter_(2, ind, val)
+    return probs
+
+
+def mask_by_random_topk(mask_len, probs, temperature=1.0, generator=None):
+    confidence = log(probs) + temperature * gumbel_noise(probs, generator=generator)
+    sorted_confidence = torch.sort(confidence, dim=-1).values
+    cut_off = torch.gather(sorted_confidence, 1, mask_len.long())
+    masking = confidence < cut_off
+    return masking
+
+
+def cosine_schedule(t):
+    return torch.cos(t * math.pi * 0.5)
+
+
+def linear_schedule(t):
+    mask_ratio = 1 - t
+    mask_ratio = mask_ratio.clamp(min=1e-6, max=1.0)
+    return mask_ratio
+
+
+def pow(t, method):
+    exponent = float(method.replace("pow", ""))
+    mask_ratio = 1.0 - t**exponent
+    mask_ratio = mask_ratio.clamp(min=1e-6, max=1.0)
+    return mask_ratio
+
+
+def sigmoid_schedule(t, start=-3, end=3, tau=1.0, clip_min=1e-6):
+    for item in [t, start, end, tau]:
+        item = torch.tensor(item) if not torch.is_tensor(item) else item
+
+    # A gamma function based on sigmoid function.
+    v_start = torch.sigmoid(torch.tensor(start / tau))
+    v_end = torch.sigmoid(torch.tensor(end / tau))
+    output = torch.sigmoid((t * (end - start) + start) / tau)
+    output = (v_end - output) / (v_end - v_start)
+    return torch.clip(output, clip_min, 1.0)
+
+
+def get_mask_chedule(method, **schedule_kwargs):
+    if method == "cosine":
+        return cosine_schedule
+    elif method == "linear":
+        return linear_schedule
+    elif "pow" in method:
+        return partial(pow, method=method)
+    elif method == "sigmoid":
+        return partial(sigmoid_schedule, **schedule_kwargs)
+    else:
+        raise ValueError("Unknown schedule method: {}".format(method))
+
+def top_k_top_p_filtering(
+    logits: torch.Tensor,
+    top_k: int = 0,
+    top_p: float = 1.0,
+    filter_value: float = -float("Inf"),
+    min_tokens_to_keep: int = 1,
+) -> torch.Tensor:
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+    return logits

models/training_utils.py

@@ -0,0 +1,455 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import os
+import random
+from typing import Any, Dict, Iterable, Optional, Union
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+
+
+def enable_full_determinism(seed: int):
+    """
+    Helper function for reproducible behavior during distributed training. See
+    - https://pytorch.org/docs/stable/notes/randomness.html for pytorch
+    """
+    # set seed first
+    set_seed(seed)
+
+    #  Enable PyTorch deterministic mode. This potentially requires either the environment
+    #  variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set,
+    # depending on the CUDA version, so we set them both here
+    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
+    torch.use_deterministic_algorithms(True)
+
+    # Enable CUDNN deterministic mode
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def set_seed(seed: int):
+    """
+    Args:
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+        seed (`int`): The seed to set.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    # ^^ safe to call this function even if cuda is not available
+
+
+# Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14
+class EMA:
+    """
+    Exponential Moving Average of models weights
+    """
+
+    def __init__(
+        self,
+        parameters: Iterable[torch.nn.Parameter],
+        decay: float = 0.9999,
+        min_decay: float = 0.0,
+        update_after_step: int = 0,
+        use_ema_warmup: bool = False,
+        inv_gamma: Union[float, int] = 1.0,
+        power: Union[float, int] = 2 / 3,
+        model_cls: Optional[Any] = None,
+        model_config: Dict[str, Any] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            parameters (Iterable[torch.nn.Parameter]): The parameters to track.
+            decay (float): The decay factor for the exponential moving average.
+            min_decay (float): The minimum decay factor for the exponential moving average.
+            update_after_step (int): The number of steps to wait before starting to update the EMA weights.
+            use_ema_warmup (bool): Whether to use EMA warmup.
+            inv_gamma (float):
+                Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
+            power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
+            device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA
+                        weights will be stored on CPU.
+
+        @crowsonkb's notes on EMA Warmup:
+            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
+            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
+            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
+            at 215.4k steps).
+        """
+
+        parameters = list(parameters)
+        self.shadow_params = [p.clone().detach() for p in parameters]
+
+        self.temp_stored_params = None
+
+        self.decay = decay
+        self.min_decay = min_decay
+        self.update_after_step = update_after_step
+        self.use_ema_warmup = use_ema_warmup
+        self.inv_gamma = inv_gamma
+        self.power = power
+        self.optimization_step = 0
+        self.cur_decay_value = None  # set in `step()`
+
+        self.model_cls = model_cls
+        self.model_config = model_config
+
+    @classmethod
+    def from_pretrained(cls, path, model_cls) -> "EMA":
+        _, ema_kwargs = model_cls.load_config(path, return_unused_kwargs=True)
+        model = model_cls.from_pretrained(path)
+
+        ema_model = cls(model.parameters(), model_cls=model_cls, model_config=model.config)
+
+        ema_model.load_state_dict(ema_kwargs)
+        return ema_model
+
+    def save_pretrained(self, path):
+        if self.model_cls is None:
+            raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")
+
+        if self.model_config is None:
+            raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")
+
+        model = self.model_cls.from_config(self.model_config)
+        state_dict = self.state_dict()
+        state_dict.pop("shadow_params", None)
+
+        model.register_to_config(**state_dict)
+        self.copy_to(model.parameters())
+        model.save_pretrained(path)
+
+    def get_decay(self, optimization_step: int) -> float:
+        """
+        Compute the decay factor for the exponential moving average.
+        """
+        step = max(0, optimization_step - self.update_after_step - 1)
+
+        if step <= 0:
+            return 0.0
+
+        if self.use_ema_warmup:
+            cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power
+        else:
+            cur_decay_value = (1 + step) / (10 + step)
+
+        cur_decay_value = min(cur_decay_value, self.decay)
+        # make sure decay is not smaller than min_decay
+        cur_decay_value = max(cur_decay_value, self.min_decay)
+        return cur_decay_value
+
+    @torch.no_grad()
+    def step(self, parameters: Iterable[torch.nn.Parameter]):
+        parameters = list(parameters)
+
+        self.optimization_step += 1
+
+        # Compute the decay factor for the exponential moving average.
+        decay = self.get_decay(self.optimization_step)
+        self.cur_decay_value = decay
+        one_minus_decay = 1 - decay
+
+        for s_param, param in zip(self.shadow_params, parameters):
+            if param.requires_grad:
+                s_param.sub_(one_minus_decay * (s_param - param))
+            else:
+                s_param.copy_(param)
+
+        torch.cuda.empty_cache()
+
+    def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        """
+        Copy current averaged parameters into given collection of parameters.
+
+        Args:
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored moving averages. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        parameters = list(parameters)
+        for s_param, param in zip(self.shadow_params, parameters):
+            param.data.copy_(s_param.to(param.device).data)
+
+    def to(self, device=None, dtype=None) -> None:
+        r"""Move internal buffers of the ExponentialMovingAverage to `device`.
+
+        Args:
+            device: like `device` argument to `torch.Tensor.to`
+        """
+        # .to() on the tensors handles None correctly
+        self.shadow_params = [
+            p.to(device=device, dtype=dtype) if p.is_floating_point() else p.to(device=device)
+            for p in self.shadow_params
+        ]
+
+    def state_dict(self) -> dict:
+        r"""
+        Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during
+        checkpointing to save the ema state dict.
+        """
+        # Following PyTorch conventions, references to tensors are returned:
+        # "returns a reference to the state and not its copy!" -
+        # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict
+        return {
+            "decay": self.decay,
+            "min_decay": self.min_decay,
+            "optimization_step": self.optimization_step,
+            "update_after_step": self.update_after_step,
+            "use_ema_warmup": self.use_ema_warmup,
+            "inv_gamma": self.inv_gamma,
+            "power": self.power,
+            "shadow_params": self.shadow_params,
+        }
+
+    def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""
+        Args:
+        Save the current parameters for restoring later.
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                temporarily stored.
+        """
+        self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]
+
+    def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""
+        Args:
+        Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without:
+        affecting the original optimization process. Store the parameters before the `copy_to()` method. After
+        validation (or model saving), use this to restore the former parameters.
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored parameters. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        if self.temp_stored_params is None:
+            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights to `restore()`")
+        for c_param, param in zip(self.temp_stored_params, parameters):
+            param.data.copy_(c_param.data)
+
+        # Better memory-wise.
+        self.temp_stored_params = None
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        r"""
+        Args:
+        Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the
+        ema state dict.
+            state_dict (dict): EMA state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        # deepcopy, to be consistent with module API
+        state_dict = copy.deepcopy(state_dict)
+
+        self.decay = state_dict.get("decay", self.decay)
+        if self.decay < 0.0 or self.decay > 1.0:
+            raise ValueError("Decay must be between 0 and 1")
+
+        self.min_decay = state_dict.get("min_decay", self.min_decay)
+        if not isinstance(self.min_decay, float):
+            raise ValueError("Invalid min_decay")
+
+        self.optimization_step = state_dict.get("optimization_step", self.optimization_step)
+        if not isinstance(self.optimization_step, int):
+            raise ValueError("Invalid optimization_step")
+
+        self.update_after_step = state_dict.get("update_after_step", self.update_after_step)
+        if not isinstance(self.update_after_step, int):
+            raise ValueError("Invalid update_after_step")
+
+        self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)
+        if not isinstance(self.use_ema_warmup, bool):
+            raise ValueError("Invalid use_ema_warmup")
+
+        self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)
+        if not isinstance(self.inv_gamma, (float, int)):
+            raise ValueError("Invalid inv_gamma")
+
+        self.power = state_dict.get("power", self.power)
+        if not isinstance(self.power, (float, int)):
+            raise ValueError("Invalid power")
+
+        shadow_params = state_dict.get("shadow_params", None)
+        if shadow_params is not None:
+            self.shadow_params = shadow_params
+            if not isinstance(self.shadow_params, list):
+                raise ValueError("shadow_params must be a list")
+            if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):
+                raise ValueError("shadow_params must all be Tensors")
+
+
+# calculates entropy over each pixel distribution
+def pixel_entropy_per_percent_masked_bucket(logits, input_ids, mask_id):
+    # only calculated entropy over image tokens that were masked in the original image
+    masked_tokens = input_ids == mask_id
+    num_masked_pixels = masked_tokens.sum(-1)
+
+    probs = F.softmax(logits, dim=-1)
+    log_probs = F.log_softmax(logits, dim=-1)
+
+    entropy_per_pixel = -((probs * log_probs).sum(-1))
+
+    # the predictions for non-masked aren't used, so set their entropies to zero
+    entropy_per_pixel[~masked_tokens] = 0
+
+    entropy_per_image_numerator = entropy_per_pixel.sum(-1)
+    entropy_per_image = entropy_per_image_numerator / num_masked_pixels
+
+    total_buckets = 10
+    masked_buckets = input_ids_to_masked_buckets(input_ids, mask_id, total_buckets)
+
+    entropy_by_masked_bucket = average_by_buckets(entropy_per_image, masked_buckets, total_buckets)
+
+    return entropy_by_masked_bucket
+
+
+# calculates entropy over the averaged distribution of pixels for the whole image
+def image_entropy_per_percent_masked_bucket(logits, input_ids, mask_id):
+    # only calculated entropy over image tokens that were masked in the original image
+    masked_tokens = input_ids == mask_id
+    num_masked_pixels = masked_tokens.sum(-1, keepdim=True)
+
+    pixel_probs = F.softmax(logits, dim=-1)
+    pixel_probs[~masked_tokens] = 0
+    image_probs_numerator = pixel_probs.sum(-2)
+    image_probs = image_probs_numerator / num_masked_pixels
+
+    image_log_probs = image_probs.log()
+
+    entropy_per_image = -((image_probs * image_log_probs).sum(-1))
+
+    total_buckets = 10
+    masked_buckets = input_ids_to_masked_buckets(input_ids, mask_id, total_buckets)
+
+    entropy_by_masked_bucket = average_by_buckets(entropy_per_image, masked_buckets, total_buckets)
+
+    return entropy_by_masked_bucket
+
+
+def cross_entropy_per_percent_masked_bucket(logits, labels, input_ids, mask_id, output_size, label_smoothing):
+    cross_entropy_per_image = F.cross_entropy(
+        logits.view(-1, output_size),
+        labels.view(-1),
+        ignore_index=-100,
+        label_smoothing=label_smoothing,
+        reduction="none",
+    )
+
+    total_buckets = 10
+    masked_buckets = input_ids_to_masked_buckets(input_ids, mask_id, total_buckets)
+
+    cross_entropy_by_percent_masked_bucket = average_by_buckets(cross_entropy_per_image, masked_buckets, total_buckets)
+
+    return cross_entropy_by_percent_masked_bucket
+
+
+def token_probability_distributions_per_percent_masked_bucket(logits, input_ids, mask_id):
+    probs = F.softmax(logits, dim=-1)
+
+    total_buckets = 10
+    masked_buckets = input_ids_to_masked_buckets(input_ids, mask_id, total_buckets)
+
+    data = []
+
+    for bucket_idx in range(total_buckets):
+        indices_for_bucket = masked_buckets[masked_buckets == bucket_idx]
+
+        # It's ok if none were noised in the range of this bucket. This
+        # function will be called for a later training step where it's likely
+        # there will be an element noised in the range.
+        if indices_for_bucket.shape[0] == 0:
+            continue
+
+        index_for_bucket = indices_for_bucket[0]
+
+        image_probs = probs[index_for_bucket]
+
+        # find the index of a masked pixel for the image
+        input_ids_for_image = input_ids[index_for_bucket]
+        masked_pixels_probs = image_probs[input_ids_for_image == mask_id]
+
+        masked_pixel_probs = masked_pixels_probs[0]
+
+        masked_pixel_probs = masked_pixel_probs.cpu().numpy()
+
+        for masked_pixel_prob in masked_pixel_probs:
+            data.append({"bucket": bucket_idx, "masked_pixel_prob": masked_pixel_prob})
+
+    df = pd.DataFrame(data)
+
+    return df
+
+
+def average_by_buckets(values, masked_buckets, total_buckets):
+    unique_buckets, bucket_counts = masked_buckets.unique(dim=0, return_counts=True)
+
+    numerator = torch.zeros(total_buckets, device=values.device)
+
+    numerator.scatter_add_(0, masked_buckets, values)
+
+    # default value is one because the buckets for which there aren't
+    # any values will have a numerator of zero. So we just need to not divide
+    # by zero.
+    denominator = torch.ones(total_buckets, device=values.device, dtype=torch.long)
+    denominator[unique_buckets] = bucket_counts
+
+    averaged_by_buckets = numerator / denominator
+
+    return averaged_by_buckets
+
+
+def input_ids_to_masked_buckets(input_ids, mask_id, total_buckets=10):
+    assert total_buckets == 10
+
+    masked_percent = (input_ids == mask_id).sum(-1) / input_ids.shape[-1]
+
+    # we do not formally use timesteps to noise images. Instead, we mask a percent
+    # of the pixels. We don't want to log entropy for every mask percent between 0 and 1,
+    # and we also want to track how the entropy evolves over time w/in a range of mask
+    # percents that should have similar entropy. So we bucket the masked percents into a
+    # fixed number of buckets
+
+    # we could generalize this later if needed but for now, let's just assume a fixed
+    # number of 10 buckets.
+
+    # How this maps to a bucket index:
+    # (mask) * bucket_index +
+    # (mask_1) * bucket_index_1
+    #
+    # -> Where the mask is true will be set to the expected bucket index,
+    # where the mask is false will be set to 0.
+    #
+    # Given the probabilities are between 0 and 1, each masked_percent will get mapped
+    # to a timestep by one and only one of the masks.
+
+    masked_buckets = (
+        ((0 < masked_percent) & (masked_percent <= 0.1)) * 0
+        + ((0.1 < masked_percent) & (masked_percent <= 0.2)) * 1
+        + ((0.2 < masked_percent) & (masked_percent <= 0.3)) * 2
+        + ((0.3 < masked_percent) & (masked_percent <= 0.4)) * 3
+        + ((0.4 < masked_percent) & (masked_percent <= 0.5)) * 4
+        + ((0.5 < masked_percent) & (masked_percent <= 0.6)) * 5
+        + ((0.6 < masked_percent) & (masked_percent <= 0.7)) * 6
+        + ((0.7 < masked_percent) & (masked_percent <= 0.8)) * 7
+        + ((0.8 < masked_percent) & (masked_percent <= 0.9)) * 8
+        + ((0.9 < masked_percent) & (masked_percent <= 1.0)) * 9
+    )
+
+    return masked_buckets

prompting_utils.py

@@ -0,0 +1,528 @@
+import torch
+
+class UniversalPrompting():
+    def __init__(self, text_tokenizer,
+                 special_tokens=("<|soi|>", "<|eoi|>", "<|sov|>", "<|eov|>", "<|t2i|>", "<|mmu|>", "<|t2v|>", "<|v2v|>", "<|lvg|>"),
+                 max_text_len=8000, max_seq_len=377, ignore_id=-100, cond_dropout_prob=0.1):
+        """
+        :param text_tokenizer: original text tokenizer
+        """
+        self.text_tokenizer = text_tokenizer
+        self.text_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        self.text_tokenizer.add_tokens(list(special_tokens))
+        self.sptids_dict = {token: torch.tensor(self.text_tokenizer.convert_tokens_to_ids([token])) for token in
+                            special_tokens}
+        self.sptids_dict['<|sot|>'] = torch.tensor([self.text_tokenizer.bos_token_id])
+        self.sptids_dict['<|eot|>'] = torch.tensor([self.text_tokenizer.eos_token_id])
+        self.sptids_dict['<|pad|>'] = torch.tensor([self.text_tokenizer.pad_token_id])
+        # plus 1 because at this time we add a task token before
+        self.max_text_len = max_text_len + 1
+        self.pad_id = self.text_tokenizer.convert_tokens_to_ids('[PAD]')
+        self.ignore_id = ignore_id
+        self.cond_dropout_prob = cond_dropout_prob
+
+    def t2i_prompt_predict_next(self, text_ids, image_ids, labels):
+
+        device = image_ids.device
+        sequence_ids = []
+        attention_masks = []
+        label_ids = []
+        probs = torch.rand(len(text_ids))
+        for i in range(len(text_ids)):
+
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.bos_token_id] + text_ids[i]
+
+            temp_ids = [int(self.sptids_dict['<|t2i|>'])] + text_ids[i] + [self.text_tokenizer.eos_token_id]
+
+            # randomly dropout text condition
+            if probs[i] < self.cond_dropout_prob:
+                temp_ids = [int(self.sptids_dict['<|t2i|>']), self.text_tokenizer.bos_token_id, self.text_tokenizer.eos_token_id]
+
+            if self.max_text_len >= len(temp_ids):
+                temp_ids = [self.pad_id] * (self.max_text_len - len(temp_ids)) + temp_ids
+                temp_masks = [0] * (self.max_text_len - len(temp_ids)) + [1] * (len(temp_ids) + image_ids.shape[-1] + 3)
+            else:
+                # should add the eos token
+                temp_ids = temp_ids[:self.max_text_len - 1] + [self.text_tokenizer.eos_token_id]
+                temp_masks = [1] * (len(temp_ids) + image_ids.shape[-1] + 3)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_label_ids = torch.cat([
+                # should we predict text tokens when doing image reconstruction?
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                labels[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_label_ids = torch.where(temp_label_ids == self.pad_id, self.ignore_id, temp_label_ids)
+
+            temp_ids = torch.cat([
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                image_ids[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_masks = torch.tensor(temp_masks).to(device)
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+            label_ids.append(temp_label_ids.unsqueeze(0))
+
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0), torch.cat(label_ids, dim=0)
+
+    def t2i_gen_prompt(self, text_ids, image_ids):
+
+        device = image_ids.device
+        sequence_ids = []
+        attention_masks = []
+        for i in range(len(text_ids)):
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.bos_token_id] + text_ids[i]
+            # note that, llama3 tokenizer automatically add the bot token at first but without eot
+            temp_ids = [int(self.sptids_dict['<|t2i|>'])] + text_ids[i] + [self.text_tokenizer.eos_token_id]
+            if self.max_text_len >= len(temp_ids):
+                temp_ids = [self.pad_id] * (self.max_text_len - len(temp_ids)) + temp_ids
+                temp_masks = [0] * (self.max_text_len - len(temp_ids)) + [1] * len(temp_ids)
+            else:
+                temp_ids = temp_ids[:self.max_text_len - 1] + [self.text_tokenizer.eos_token_id]
+                temp_masks = [1] * len(temp_ids)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_ids = torch.cat([
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                image_ids[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_masks = torch.tensor(temp_masks).to(device)
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0)
+
+    # language modeling
+    def lm_prompt(self, text_ids, max_seq_len):
+
+        sequence_ids = []
+        attention_masks = []
+        label_ids = []
+        for i in range(len(text_ids)):
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.eos_token_id] + text_ids[i]
+
+            temp_ids = text_ids[i] + [self.text_tokenizer.eos_token_id]
+
+            if max_seq_len >= len(temp_ids):
+                temp_labels_ids = temp_ids + [self.ignore_id] * (max_seq_len - len(temp_ids))
+                temp_ids = temp_ids + [self.pad_id] * (max_seq_len - len(temp_ids))
+                temp_masks = [1] * len(temp_ids) + [0] * (max_seq_len - len(temp_ids))
+            else:
+                # In language modeling, we only process text tokens. We do not add the eos token if the text length
+                # exceeds the max sequence length
+                temp_labels_ids = temp_ids[:max_seq_len]
+                temp_ids = temp_ids[:max_seq_len]
+                temp_masks = [1] * len(temp_ids)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_ids = torch.tensor(temp_ids)
+            temp_masks = torch.tensor(temp_masks)
+            temp_labels_ids = torch.tensor(temp_labels_ids)
+
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+            label_ids.append(temp_labels_ids.unsqueeze(0))
+
+        # input_ids, masks, labels
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0), torch.cat(label_ids, dim=0)
+
+    def mmu_prompt(self, image_ids, text_ids):
+        device = image_ids.device
+        sequence_ids = []
+        attention_masks = []
+        label_ids = []
+        max_text_len = self.max_text_len - 1
+        for i in range(len(text_ids)):
+            # note that, llama3 tokenizer automatically add the bot token at first but without eot
+            # for empty list []
+
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.eos_token_id] + text_ids[i]
+
+            temp_ids = text_ids[i] + [self.text_tokenizer.eos_token_id]
+
+            if max_text_len >= len(temp_ids):
+                # minus 1 because task token was prepended to the former image tokens
+                temp_ids = temp_ids + [self.pad_id] * (max_text_len - len(temp_ids))
+                temp_masks = [1] * (len(temp_ids) + image_ids.shape[-1] + 3) + [0] * (max_text_len - len(temp_ids))
+            else:
+                # should add the eos token
+                temp_ids = temp_ids[:max_text_len - 1] + [self.text_tokenizer.eos_token_id]
+                temp_masks = [1] * (len(temp_ids) + image_ids.shape[-1] + 3)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_label_ids = torch.cat([
+                torch.tensor([self.ignore_id]).to(device),
+                torch.tensor([self.ignore_id]).to(device),
+                torch.ones_like(image_ids[i]) * self.ignore_id,
+                torch.tensor([self.ignore_id]).to(device),
+                torch.tensor(temp_ids).to(device),
+            ], dim=0)
+
+            temp_label_ids = torch.where(temp_label_ids == self.pad_id, self.ignore_id, temp_label_ids)
+
+            temp_ids = torch.cat([
+                self.sptids_dict['<|mmu|>'].to(device),  # task token
+                self.sptids_dict['<|soi|>'].to(device),
+                image_ids[i],
+                self.sptids_dict['<|eoi|>'].to(device),
+                torch.tensor(temp_ids).to(device),
+            ], dim=0)
+
+            temp_masks = torch.tensor(temp_masks).to(device)
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+            label_ids.append(temp_label_ids.unsqueeze(0))
+
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0), torch.cat(label_ids, dim=0)
+
+    def t2v_prompt(self, text_ids, video_ids):
+        """
+        :param text_ids:
+        :param video_ids:
+        :return:
+        """
+        pass
+
+    def i2v_prompt(self, image_ids, video_ids):
+        """
+        :param image_ids:
+        :param video_ids:
+        :return:
+        """
+        pass
+
+    def lvg_prompt(self, text_ids, image_ids, labels):
+
+        device = image_ids.device
+        sequence_ids = []
+        attention_masks = []
+        label_ids = []
+        probs = torch.rand(len(text_ids))
+        probs2 = torch.rand(len(text_ids))
+        for i in range(len(text_ids)):
+
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.bos_token_id] + text_ids[i]
+
+            temp_ids = [int(self.sptids_dict['<|t2i|>'])] + text_ids[i] + [self.text_tokenizer.eos_token_id]
+
+            # randomly dropout text condition
+            if probs[i] < self.cond_dropout_prob:
+                temp_ids = [int(self.sptids_dict['<|t2i|>']), self.text_tokenizer.bos_token_id,
+                            self.text_tokenizer.eos_token_id]
+
+            if self.max_text_len >= len(temp_ids):
+                temp_ids = [self.pad_id] * (self.max_text_len - len(temp_ids)) + temp_ids
+                temp_masks = [0] * (self.max_text_len - len(temp_ids)) + [1] * (len(temp_ids) + image_ids.shape[-1] + 3)
+            else:
+                # should add the eos token
+                temp_ids = temp_ids[:self.max_text_len - 1] + [self.text_tokenizer.eos_token_id]
+                temp_masks = [1] * (len(temp_ids) + image_ids.shape[-1] + 3)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_label_ids = torch.cat([
+                # should we predict text tokens when doing image reconstruction?
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                labels[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_label_ids = torch.where(temp_label_ids == self.pad_id, self.ignore_id, temp_label_ids)
+
+            temp_ids = torch.cat([
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                image_ids[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_masks = torch.tensor(temp_masks).to(device)
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+            label_ids.append(temp_label_ids.unsqueeze(0))
+
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0), torch.cat(label_ids, dim=0)
+
+    def lvg_gen_prompt(self, text_ids, image_ids):
+
+        device = image_ids.device
+        sequence_ids = []
+        attention_masks = []
+        for i in range(len(text_ids)):
+            if len(text_ids[i]) == 0:
+                text_ids[i] = [self.text_tokenizer.bos_token_id]
+            elif text_ids[i][0] != self.text_tokenizer.bos_token_id:
+                text_ids[i] = [self.text_tokenizer.bos_token_id] + text_ids[i]
+            # note that, llama3 tokenizer automatically add the bot token at first but without eot
+            temp_ids = [int(self.sptids_dict['<|t2i|>'])] + text_ids[i] + [self.text_tokenizer.eos_token_id]
+            if self.max_text_len >= len(temp_ids):
+                temp_ids = [self.pad_id] * (self.max_text_len - len(temp_ids)) + temp_ids
+                temp_masks = [0] * (self.max_text_len - len(temp_ids)) + [1] * len(temp_ids)
+            else:
+                temp_ids = temp_ids[:self.max_text_len - 1] + [self.text_tokenizer.eos_token_id]
+                temp_masks = [1] * len(temp_ids)  # +2 for two special tokens
+
+            # prompting -- [task token] [sot] [text tokens] [eot] [soi] [image tokens] [eoi]
+            temp_ids = torch.cat([
+                torch.tensor(temp_ids).to(device),
+                self.sptids_dict['<|soi|>'].to(device),
+                image_ids[i],
+                self.sptids_dict['<|eoi|>'].to(device)
+            ], dim=0)
+
+            temp_masks = torch.tensor(temp_masks).to(device)
+            sequence_ids.append(temp_ids.unsqueeze(0))
+            attention_masks.append(temp_masks.unsqueeze(0))
+
+        return torch.cat(sequence_ids, dim=0), torch.cat(attention_masks, dim=0)
+
+    def mask_prompt(self):
+        pass
+
+    def __call__(self, input, task, padding=True, config=None):
+        """
+        input (tuple) : data pairs contain text(str), image(tensor), or videos(tensor).
+        task (str) : a flag indicates the current task.
+        """
+        if task == "t2i":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.t2i_prompt(text_ids, image_ids, input[2])
+
+        elif task == "t2i_predict_next":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.t2i_prompt_predict_next(text_ids, image_ids, input[2])
+
+        elif task == "t2i_predict_next_plus_lm":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.t2i_prompt_predict_next(text_ids[:config.training.batch_size], image_ids,
+                                                                   input[2])
+            sequence_ids_with_masks_lm = self.lm_prompt(text_ids[config.training.batch_size:], input[3])
+            return sequence_ids_with_masks, sequence_ids_with_masks_lm
+
+        elif task == "t2i_gen":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.t2i_gen_prompt(text_ids, image_ids)
+
+        elif task == "lm":
+            text_ids = self.text_tokenizer(input[0], truncation=True)['input_ids']  # (B, max_len)
+            sequence_ids_with_masks = self.lm_prompt(text_ids, input[1])
+
+        elif task == "mmu":
+            image_ids = input[0]
+            text_ids = self.text_tokenizer(input[1])['input_ids']
+            sequence_ids_with_masks = self.mmu_prompt(image_ids, text_ids)
+
+        elif task == "t2v":
+            text_ids = self.text_tokenizer(input[0]['input_ids'])
+            video_ids = self.vision_tokenizer(input[1])
+            sequence_ids_with_masks = self.t2v_prompt(text_ids, video_ids)
+
+        elif task == "i2v":
+            image_ids = self.text_tokenizer(input[0])
+            video_ids = self.vision_tokenizer(input[1])
+            sequence_ids_with_masks = self.i2v_prompt(image_ids, video_ids)
+
+        elif task == "lvg":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.lvg_prompt(text_ids, image_ids, input[2])
+
+        elif task == "lvg_gen":
+            text_ids = self.text_tokenizer(input[0])['input_ids']  # (B, max_len)
+            image_ids = input[1]  # (B, #tokens)
+            sequence_ids_with_masks = self.lvg_gen_prompt(text_ids, image_ids)
+        else:
+            raise NotImplementedError
+
+        return sequence_ids_with_masks
+
+def create_attention_mask_predict_next(sequence, pad_id=128256, soi_id=128257, eoi_id=128258, rm_pad_in_image=False,
+                                       return_inverse_mask=True):
+    # sequence is expected to be of shape [N, L]
+    N, L = sequence.shape
+
+    # Masks to identify different types of tokens
+    is_padding = sequence == pad_id
+
+    is_start_image = sequence == soi_id
+
+    is_end_image = sequence == eoi_id
+
+    # Create cumulative sum masks to identify regions of image tokens
+    cumulative_start = torch.cumsum(is_start_image, dim=1)
+    cumulative_end = torch.cumsum(is_end_image, dim=1)
+    in_image_segment = (cumulative_start > cumulative_end) | is_start_image | is_end_image
+
+    is_text = ~(in_image_segment)
+
+    causal_mask = torch.tril(torch.ones((L, L), dtype=torch.bool)).to(sequence.device)
+
+    mask_text = is_text[:, :, None] * causal_mask[None, :, :]
+
+    is_text_image = is_text | in_image_segment
+
+    mask_text_image_bi = is_text_image[:, :, None] * is_text_image[:, None, :]
+    if rm_pad_in_image:
+        sid_img = torch.where(sequence == soi_id)[1]
+        for i in range(mask_text_image_bi.shape[0]):
+            pad_end_idx = torch.where(sequence[i] == pad_id)
+            if len(pad_end_idx[0]) != 0:
+                pad_end_idx = pad_end_idx[0][-1]
+                mask_text[i][pad_end_idx + 1:, :pad_end_idx + 1] = 0
+            id_padding = torch.where(is_padding[i] == True)
+            mask_text_image_bi[i][sid_img[i]:, id_padding[0]] = 0
+
+    mask_text[in_image_segment] = mask_text_image_bi[in_image_segment]
+    # No token attends to padding tokens and padding tokens do not attend to any token
+    if return_inverse_mask:
+        inverted_mask = 1.0 - mask_text.type(sequence.dtype)
+        inverted_mask = inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.iinfo(sequence.dtype).min
+        )
+        return inverted_mask.unsqueeze(1)
+    else:
+        return mask_text.unsqueeze(1)
+
+def create_attention_mask_lvg(sequence, pad_id=128256, soi_id=128257, eoi_id=128258, return_inverse_mask=True):
+    # sequence is expected to be of shape [N, L]
+    N, L = sequence.shape
+    # Masks to identify different types of tokens
+    is_padding = sequence == pad_id
+    mask_text_image_bi = torch.tril(torch.ones(N, L, L), diagonal=0).to(sequence.device)
+
+    sid_img = torch.where(sequence == soi_id)[1].reshape(mask_text_image_bi.shape[0], -1)[:, 0]
+    sid_img_for_bi = torch.where(sequence == soi_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    eid_img_for_bi = torch.where(sequence == eoi_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    for i in range(N):
+        id_padding = torch.where(is_padding[i] == True)
+        mask_text_image_bi[i][sid_img[i]:, id_padding[0]] = 0
+        for j in range(sid_img_for_bi.shape[-1]):
+            mask_text_image_bi[i][sid_img_for_bi[i, j]:eid_img_for_bi[i, j] + 1,
+            sid_img_for_bi[i, j]:eid_img_for_bi[i, j] + 1] = 1
+
+    # No token attends to padding tokens and padding tokens do not attend to any token
+    if return_inverse_mask:
+        inverted_mask = 1.0 - mask_text_image_bi.type(sequence.dtype)
+        inverted_mask = inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.iinfo(sequence.dtype).min
+        )
+        return inverted_mask.unsqueeze(1)
+    else:
+        return mask_text_image_bi.unsqueeze(1)
+
+# texts without attending image regions
+def create_attention_mask_lvg_v2(sequence, pad_id=128256, soi_id=128257, eoi_id=128258, sot_id=1000, eot_id=1001, return_inverse_mask=True):
+    # sequence is expected to be of shape [N, L]
+    N, L = sequence.shape
+    # Masks to identify different types of tokens
+    is_padding = sequence == pad_id
+    # is_text = torch.where(sequence < 2000, True, False)
+    is_text = torch.where(sequence < pad_id, True, False)
+    mask_text_image_bi = torch.tril(torch.ones(N, L, L), diagonal=0).to(sequence.device).int()
+    sid_text_for_bi = torch.where(sequence == sot_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    eid_text_for_bi = torch.where(sequence == eot_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    # import ipdb
+    # ipdb.set_trace()
+    if sot_id == eot_id:
+        if sid_text_for_bi.shape[-1] % 2 != 0:
+            sid_text_for_bi = sid_text_for_bi[:, :-1]
+            eid_text_for_bi = eid_text_for_bi[:, :-1]
+        select_idx = [i for i in range(0, sid_text_for_bi.shape[1], 2)]
+        sid_text_for_bi = sid_text_for_bi[:, select_idx]
+        select_idx = [i+1 for i in range(0, eid_text_for_bi.shape[1], 2)]
+        eid_text_for_bi = eid_text_for_bi[:, select_idx]
+    sid_img_for_bi = torch.where(sequence == soi_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    eid_img_for_bi = torch.where(sequence == eoi_id)[1].reshape(mask_text_image_bi.shape[0], -1)
+    all_zeros = torch.zeros_like(mask_text_image_bi).int()
+    for i in range(N):
+        all_zeros[i, :, is_text[i]] = 1
+        for j in range(sid_text_for_bi.shape[-1]):
+            all_zeros[i][is_text[i], sid_text_for_bi[i, j]:eid_text_for_bi[i, j]+1] = 1
+            all_zeros[i][~is_text[i], sid_text_for_bi[i, j]:eid_text_for_bi[i, j]+1] = 1
+        for j in range(sid_img_for_bi.shape[-1]):
+            all_zeros[i][~is_text[i], sid_img_for_bi[i, j]:eid_img_for_bi[i, j]+1] = 1
+    mask_text_image_bi = mask_text_image_bi * all_zeros
+    sid_img = torch.where(sequence == soi_id)[1].reshape(mask_text_image_bi.shape[0], -1)[:, 0]
+
+    for i in range(N):
+        id_padding = torch.where(is_padding[i] == True)
+        mask_text_image_bi[i][sid_img[i]:, id_padding[0]] = 0
+        for j in range(sid_img_for_bi.shape[-1]):
+            mask_text_image_bi[i][sid_img_for_bi[i, j]:eid_img_for_bi[i, j]+1, sid_img_for_bi[i, j]:eid_img_for_bi[i, j]+1] = 1
+
+    mask_text_image_bi[:, :, 0] = 1
+    # No token attends to padding tokens and padding tokens do not attend to any token
+    if return_inverse_mask:
+        inverted_mask = 1.0 - mask_text_image_bi.type(sequence.dtype)
+        inverted_mask = inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.iinfo(sequence.dtype).min
+        )
+        return inverted_mask.unsqueeze(1)
+    else:
+        return mask_text_image_bi.unsqueeze(1)
+
+def create_attention_mask_for_mmu(sequence, eoi_id=128258, return_inverse_mask=True):
+    N, L = sequence.shape
+    causal_mask = torch.tril(torch.ones((N, 1, L, L), dtype=torch.bool)).to(sequence.device)
+    eoi_image = torch.where(sequence == eoi_id)[1]
+    causal_mask[:, :, :, :eoi_image[0] + 1] = 1
+
+    if return_inverse_mask:
+        inverted_mask = 1.0 - causal_mask.type(sequence.dtype)
+        inverted_mask = inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.iinfo(sequence.dtype).min
+        )
+        return inverted_mask
+    else:
+        return causal_mask
+
+def create_attention_mask_for_mmu_vit(
+        sequence,
+        return_inverse_mask=True,
+        system_prompt_len=0
+):
+    N, L, H = sequence.shape
+    causal_mask = torch.tril(torch.ones((N, 1, L, L), dtype=torch.bool)).to(sequence.device)
+    index = 1 + system_prompt_len + 1 + 576
+
+    causal_mask[:, :, :, :index] = 1
+    if return_inverse_mask:
+        inverted_mask = 1.0 - causal_mask.type(torch.int64)
+        inverted_mask = inverted_mask.masked_fill(
+            inverted_mask.to(torch.bool), torch.iinfo(torch.int64).min
+        )
+        return inverted_mask
+    else:
+        return causal_mask
+
+
+if __name__ == '__main__':
+    pass

requirements.txt

@@ -0,0 +1,228 @@
+accelerate==0.21.0
+aiohttp==3.9.5
+aiosignal==1.3.1
+albumentations==0.3.2
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anykeystore==0.2
+asn1crypto==1.5.1
+asttokens==2.4.1
+async-timeout==4.0.3
+attrs==21.2.0
+bidict==0.23.1
+blessed==1.20.0
+boto3==1.34.113
+botocore==1.34.113
+braceexpand==0.1.7
+cachetools==5.3.3
+certifi==2024.2.2
+cffi==1.16.0
+chardet==5.2.0
+charset-normalizer==3.3.2
+click==8.1.7
+clip==0.2.0
+clip-openai==1.0.post20230121
+cmake==3.29.3
+cramjam==2.8.3
+crcmod==1.7
+cryptacular==1.6.2
+cryptography==39.0.2
+cycler==0.12.1
+datasets
+diffusers==0.30.1
+decorator==5.1.1
+decord==0.6.0
+deepspeed==0.14.2
+defusedxml==0.7.1
+Deprecated==1.2.14
+descartes==1.1.0
+dill==0.3.8
+distlib==0.3.8
+distro-info==1.0
+dnspython==2.6.1
+docker-pycreds==0.4.0
+docstring_parser==0.16
+ecdsa==0.19.0
+einops==0.6.0
+exceptiongroup==1.2.1
+executing==2.0.1
+fairscale==0.4.13
+fastparquet==2024.5.0
+ffmpegcv==0.3.13
+filelock==3.14.0
+fire==0.6.0
+fonttools==4.51.0
+frozenlist==1.4.1
+fsspec==2023.6.0
+ftfy==6.2.0
+gitdb==4.0.11
+GitPython==3.1.43
+gpustat==1.1.1
+greenlet==3.0.3
+grpcio==1.64.0
+h11==0.14.0
+hjson==3.1.0
+huggingface-hub==0.23.2
+hupper==1.12.1
+idna==3.7
+imageio==2.34.1
+imgaug==0.2.6
+iniconfig==2.0.0
+ipaddress==1.0.23
+ipdb==0.13.13
+ipython==8.18.1
+jaxtyping==0.2.28
+jedi==0.19.1
+Jinja2==3.1.4
+jmespath==1.0.1
+joblib==1.4.2
+jsonargparse==4.14.1
+jsonlines==4.0.0
+kiwisolver==1.4.5
+kornia==0.7.2
+kornia_rs==0.1.3
+lazy_loader==0.4
+lightning==2.2.3
+lightning-utilities==0.11.2
+lit==18.1.6
+MarkupSafe==2.1.5
+matplotlib==3.5.3
+matplotlib-inline==0.1.7
+miscreant==0.3.0
+mpmath==1.3.0
+msgpack==1.0.8
+multidict==6.0.5
+multiprocess==0.70.16
+natsort==8.4.0
+networkx==3.2.1
+ninja==1.11.1.1
+numpy==1.24.4
+nuscenes-devkit==1.1.11
+oauthlib==3.2.2
+omegaconf==2.3.0
+open-clip-torch==2.24.0
+openai-clip
+opencv-python==4.9.0.80
+opencv-python-headless==3.4.18.65
+packaging==22.0
+pandas==1.5.3
+parquet==1.3.1
+parso==0.8.4
+PasteDeploy==3.1.0
+pathlib2==2.3.7.post1
+pathtools==0.1.2
+pbkdf2==1.3
+pexpect==4.9.0
+pillow==10.3.0
+plaster==1.1.2
+plaster-pastedeploy==1.0.1
+platformdirs==4.2.2
+plotly==5.22.0
+pluggy==1.5.0
+ply==3.11
+promise==2.3
+prompt-toolkit==3.0.43
+protobuf==3.20.3
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py==1.11.0
+py-cpuinfo==9.0.0
+py-spy==0.3.14
+pyarrow==11.0.0
+pyarrow-hotfix==0.6
+pyasn1==0.6.0
+pycocotools==2.0.7
+pycparser==2.22
+pycryptodomex==3.20.0
+pycurl==7.43.0.6
+pydantic==1.10.15
+pydantic_core==2.18.3
+Pygments==2.18.0
+PyJWT==2.8.0
+pynvml==11.5.0
+pyope==0.2.2
+pyOpenSSL==23.2.0
+pyparsing==3.1.2
+pyquaternion==0.9.9
+pyramid==2.0.2
+pyramid-mailer==0.15.1
+pytest==6.2.5
+python-consul==1.1.0
+python-dateutil==2.9.0.post0
+python-engineio==4.9.1
+python-etcd==0.4.5
+python-jose==3.3.0
+python-socketio==5.11.2
+python3-openid==3.2.0
+pytorch-extension==0.2
+pytorch-lightning==2.2.3
+pytz==2024.1
+PyYAML==6.0.1
+regex==2024.5.15
+repoze.sendmail==4.4.1
+requests==2.31.0
+requests-oauthlib==2.0.0
+rsa==4.9
+s3transfer==0.10.1
+safetensors==0.4.3
+schedule==1.2.2
+scikit-image==0.22.0
+scikit-learn==1.5.0
+scipy==1.13.1
+sentencepiece==0.2.0
+sentry-sdk==2.3.1
+setproctitle==1.3.3
+Shapely==1.8.5.post1
+shortuuid==1.0.13
+simple-websocket==1.0.0
+six==1.16.0
+smmap==5.0.1
+SQLAlchemy==2.0.30
+stack-data==0.6.3
+sympy==1.12
+taming-transformers-rom1504==0.0.6
+tenacity==8.3.0
+tensorboardX==2.6.2.2
+termcolor==2.4.0
+threadpoolctl==3.5.0
+thriftpy2==0.5.0
+tifffile==2024.5.22
+timm==1.0.3
+tokenizers==0.19.1
+toml==0.10.2
+tomli==2.0.1
+torch==2.2.1
+torch-fidelity==0.3.0
+torchmetrics==1.4.0.post0
+torchvision==0.17.1
+tox==3.28.0
+tqdm==4.66.4
+traitlets==5.14.3
+transaction==4.0
+transformers==4.41.1
+translationstring==1.4
+triton==2.2.0
+typeguard==2.13.3
+typing_extensions==4.12.0
+tzdata==2024.1
+urllib3==1.26.18
+velruse==1.1.1
+venusian==3.1.0
+virtualenv==20.26.2
+wandb==0.17.0
+watchdog==4.0.1
+wcwidth==0.2.13
+webdataset==0.2.86
+WebOb==1.8.7
+websocket-client==1.8.0
+wrapt==1.16.0
+wsproto==1.2.0
+WTForms==3.1.2
+wtforms-recaptcha==0.3.2
+xformers==0.0.25
+xxhash==3.4.1
+yarl==1.9.4
+zope.deprecation==5.0
+zope.interface==6.4.post2
+zope.sqlalchemy==3.1

training/__init__.py

@@ -0,0 +1 @@
+

training/conversation.py

@@ -0,0 +1,432 @@
+# Modified from LLaVA: https://github.com/haotian-liu/LLaVA.git
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n" if len(msg) > 0 else msg
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def process_image(self, image, image_process_mode, return_pil=False, image_format='PNG', max_len=1344, min_len=672):
+        if image_process_mode == "Pad":
+            def expand2square(pil_img, background_color=(122, 116, 104)):
+                width, height = pil_img.size
+                if width == height:
+                    return pil_img
+                elif width > height:
+                    result = Image.new(pil_img.mode, (width, width), background_color)
+                    result.paste(pil_img, (0, (width - height) // 2))
+                    return result
+                else:
+                    result = Image.new(pil_img.mode, (height, height), background_color)
+                    result.paste(pil_img, ((height - width) // 2, 0))
+                    return result
+            image = expand2square(image)
+        elif image_process_mode in ["Default", "Crop"]:
+            pass
+        elif image_process_mode == "Resize":
+            image = image.resize((336, 336))
+        else:
+            raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+        if max(image.size) > max_len:
+            max_hw, min_hw = max(image.size), min(image.size)
+            aspect_ratio = max_hw / min_hw
+            shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+            longest_edge = int(shortest_edge * aspect_ratio)
+            W, H = image.size
+            if H > W:
+                H, W = longest_edge, shortest_edge
+            else:
+                H, W = shortest_edge, longest_edge
+            image = image.resize((W, H))
+        if return_pil:
+            return image
+        else:
+            buffered = BytesIO()
+            image.save(buffered, format=image_format)
+            img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+            return img_b64_str
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    image = self.process_image(image, image_process_mode, return_pil=return_pil)
+                    images.append(image)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg, image, image_process_mode = msg
+                    img_b64_str = self.process_image(
+                        image, "Default", return_pil=False,
+                        image_format='JPEG')
+                    img_str = f'<img src="data:image/jpeg;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+conv_mistral_instruct = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="",
+    sep2="</s>",
+)
+
+conv_chatml_direct = Conversation(
+    system="""<|im_start|>system
+Answer the questions.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_phi3_instruct = Conversation(
+    system="""<|system|>\nYou are a helpful AI assistant.""",
+    roles=("\n<|user|>\n", "\n<|assistant|>\n"),
+    version="phi3",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|end|>",
+)
+
+# conv_phi_v0 = Conversation(
+#     system="A chat between a curious user and an artificial intelligence assistant. "
+#            "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+#     roles=("USER", "ASSISTANT"),
+#     version="v0",
+#     messages=(),
+#     offset=0,
+#     sep_style=SeparatorStyle.TWO,
+#     sep=" ",
+#     sep2="<|endoftext|>",
+# )
+
+conv_phi_v0 = Conversation(
+    system="",
+    roles=("USER", "ASSISTANT"),
+    version="v0",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="<|endoftext|>",
+)
+
+default_conversation = conv_vicuna_v1
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+    "mistral_instruct": conv_mistral_instruct,
+    "chatml_direct": conv_chatml_direct,
+    "mistral_direct": conv_chatml_direct,
+
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+    "phi3_instruct": conv_phi3_instruct,
+    "phi1.5": conv_phi_v0,
+
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

training/utils.py

@@ -0,0 +1,185 @@
+import math
+import random
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from typing import Any, List, Tuple, Union
+
+
+##################################################
+#              config utils
+##################################################
+def get_config():
+    cli_conf = OmegaConf.from_cli()
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = OmegaConf.merge(yaml_conf, cli_conf)
+
+    return conf
+
+
+def flatten_omega_conf(cfg: Any, resolve: bool = False) -> List[Tuple[str, Any]]:
+    ret = []
+
+    def handle_dict(key: Any, value: Any, resolve: bool) -> List[Tuple[str, Any]]:
+        return [(f"{key}.{k1}", v1) for k1, v1 in flatten_omega_conf(value, resolve=resolve)]
+
+    def handle_list(key: Any, value: Any, resolve: bool) -> List[Tuple[str, Any]]:
+        return [(f"{key}.{idx}", v1) for idx, v1 in flatten_omega_conf(value, resolve=resolve)]
+
+    if isinstance(cfg, DictConfig):
+        for k, v in cfg.items_ex(resolve=resolve):
+            if isinstance(v, DictConfig):
+                ret.extend(handle_dict(k, v, resolve=resolve))
+            elif isinstance(v, ListConfig):
+                ret.extend(handle_list(k, v, resolve=resolve))
+            else:
+                ret.append((str(k), v))
+    elif isinstance(cfg, ListConfig):
+        for idx, v in enumerate(cfg._iter_ex(resolve=resolve)):
+            if isinstance(v, DictConfig):
+                ret.extend(handle_dict(idx, v, resolve=resolve))
+            elif isinstance(v, ListConfig):
+                ret.extend(handle_list(idx, v, resolve=resolve))
+            else:
+                ret.append((str(idx), v))
+    else:
+        assert False
+
+    return ret
+
+
+##################################################
+#              training utils
+##################################################
+def soft_target_cross_entropy(logits, targets, soft_targets):
+    # ignore the first token from logits and targets (class id token)
+    logits = logits[:, 1:]
+    targets = targets[:, 1:]
+
+    logits = logits[..., : soft_targets.shape[-1]]
+
+    log_probs = F.log_softmax(logits, dim=-1)
+    padding_mask = targets.eq(-100)
+
+    loss = torch.sum(-soft_targets * log_probs, dim=-1)
+    loss.masked_fill_(padding_mask, 0.0)
+
+    # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded):
+    num_active_elements = padding_mask.numel() - padding_mask.long().sum()
+    loss = loss.sum() / num_active_elements
+    return loss
+
+
+def get_loss_weight(t, mask, min_val=0.3):
+    return 1 - (1 - mask) * ((1 - t) * (1 - min_val))[:, None]
+
+
+def mask_or_random_replace_tokens(image_tokens, mask_id, config, mask_schedule, is_train=True):
+    batch_size, seq_len = image_tokens.shape
+
+    if not is_train and config.training.get("eval_mask_ratios", None):
+        mask_prob = random.choices(config.training.eval_mask_ratios, k=batch_size)
+        mask_prob = torch.tensor(mask_prob, device=image_tokens.device)
+    else:
+        # Sample a random timestep for each image
+        timesteps = torch.rand(batch_size, device=image_tokens.device)
+        # Sample a random mask probability for each image using timestep and cosine schedule
+        mask_prob = mask_schedule(timesteps)
+        mask_prob = mask_prob.clip(config.training.min_masking_rate)
+
+    # creat a random mask for each image
+    num_token_masked = (seq_len * mask_prob).round().clamp(min=1)
+
+    mask_contiguous_region_prob = config.training.get("mask_contiguous_region_prob", None)
+
+    if mask_contiguous_region_prob is None:
+        mask_contiguous_region = False
+    else:
+        mask_contiguous_region = random.random() < mask_contiguous_region_prob
+
+    if not mask_contiguous_region:
+        batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)
+        mask = batch_randperm < num_token_masked.unsqueeze(-1)
+    else:
+        resolution = int(seq_len ** 0.5)
+        mask = torch.zeros((batch_size, resolution, resolution), device=image_tokens.device)
+
+        # TODO - would be nice to vectorize
+        for batch_idx, num_token_masked_ in enumerate(num_token_masked):
+            num_token_masked_ = int(num_token_masked_.item())
+
+            # NOTE: a bit handwavy with the bounds but gets a rectangle of ~num_token_masked_
+            num_token_masked_height = random.randint(
+                math.ceil(num_token_masked_ / resolution), min(resolution, num_token_masked_)
+            )
+            num_token_masked_height = min(num_token_masked_height, resolution)
+
+            num_token_masked_width = math.ceil(num_token_masked_ / num_token_masked_height)
+            num_token_masked_width = min(num_token_masked_width, resolution)
+
+            start_idx_height = random.randint(0, resolution - num_token_masked_height)
+            start_idx_width = random.randint(0, resolution - num_token_masked_width)
+
+            mask[
+            batch_idx,
+            start_idx_height: start_idx_height + num_token_masked_height,
+            start_idx_width: start_idx_width + num_token_masked_width,
+            ] = 1
+
+        mask = mask.reshape(batch_size, seq_len)
+        mask = mask.to(torch.bool)
+
+    # mask images and create input and labels
+    if config.training.get("noise_type", "mask"):
+        input_ids = torch.where(mask, mask_id, image_tokens)
+    elif config.training.get("noise_type", "random_replace"):
+        # sample random tokens from the vocabulary
+        random_tokens = torch.randint_like(
+            image_tokens, low=0, high=config.model.codebook_size, device=image_tokens.device
+        )
+        input_ids = torch.where(mask, random_tokens, image_tokens)
+    else:
+        raise ValueError(f"noise_type {config.training.noise_type} not supported")
+
+    if (
+            config.training.get("predict_all_tokens", False)
+            or config.training.get("noise_type", "mask") == "random_replace"
+    ):
+        labels = image_tokens
+        loss_weight = get_loss_weight(mask_prob, mask.long())
+    else:
+        labels = torch.where(mask, image_tokens, -100)
+        loss_weight = None
+
+    return input_ids, labels, loss_weight, mask_prob
+
+
+##################################################
+#              misc
+##################################################
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+from torchvision import transforms
+def image_transform(image, resolution=256, normalize=True):
+    image = transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR)(image)
+    image = transforms.CenterCrop((resolution, resolution))(image)
+    image = transforms.ToTensor()(image)
+    if normalize:
+        image = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)(image)
+    return image

training_utils.py

@@ -0,0 +1,185 @@
+import math
+import random
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from typing import Any, List, Tuple, Union
+
+
+##################################################
+#              config utils
+##################################################
+def get_config():
+    cli_conf = OmegaConf.from_cli()
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = OmegaConf.merge(yaml_conf, cli_conf)
+
+    return conf
+
+
+def flatten_omega_conf(cfg: Any, resolve: bool = False) -> List[Tuple[str, Any]]:
+    ret = []
+
+    def handle_dict(key: Any, value: Any, resolve: bool) -> List[Tuple[str, Any]]:
+        return [(f"{key}.{k1}", v1) for k1, v1 in flatten_omega_conf(value, resolve=resolve)]
+
+    def handle_list(key: Any, value: Any, resolve: bool) -> List[Tuple[str, Any]]:
+        return [(f"{key}.{idx}", v1) for idx, v1 in flatten_omega_conf(value, resolve=resolve)]
+
+    if isinstance(cfg, DictConfig):
+        for k, v in cfg.items_ex(resolve=resolve):
+            if isinstance(v, DictConfig):
+                ret.extend(handle_dict(k, v, resolve=resolve))
+            elif isinstance(v, ListConfig):
+                ret.extend(handle_list(k, v, resolve=resolve))
+            else:
+                ret.append((str(k), v))
+    elif isinstance(cfg, ListConfig):
+        for idx, v in enumerate(cfg._iter_ex(resolve=resolve)):
+            if isinstance(v, DictConfig):
+                ret.extend(handle_dict(idx, v, resolve=resolve))
+            elif isinstance(v, ListConfig):
+                ret.extend(handle_list(idx, v, resolve=resolve))
+            else:
+                ret.append((str(idx), v))
+    else:
+        assert False
+
+    return ret
+
+
+##################################################
+#              training utils
+##################################################
+def soft_target_cross_entropy(logits, targets, soft_targets):
+    # ignore the first token from logits and targets (class id token)
+    logits = logits[:, 1:]
+    targets = targets[:, 1:]
+
+    logits = logits[..., : soft_targets.shape[-1]]
+
+    log_probs = F.log_softmax(logits, dim=-1)
+    padding_mask = targets.eq(-100)
+
+    loss = torch.sum(-soft_targets * log_probs, dim=-1)
+    loss.masked_fill_(padding_mask, 0.0)
+
+    # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded):
+    num_active_elements = padding_mask.numel() - padding_mask.long().sum()
+    loss = loss.sum() / num_active_elements
+    return loss
+
+
+def get_loss_weight(t, mask, min_val=0.3):
+    return 1 - (1 - mask) * ((1 - t) * (1 - min_val))[:, None]
+
+
+def mask_or_random_replace_tokens(image_tokens, mask_id, config, mask_schedule, is_train=True):
+    batch_size, seq_len = image_tokens.shape
+
+    if not is_train and config.training.get("eval_mask_ratios", None):
+        mask_prob = random.choices(config.training.eval_mask_ratios, k=batch_size)
+        mask_prob = torch.tensor(mask_prob, device=image_tokens.device)
+    else:
+        # Sample a random timestep for each image
+        timesteps = torch.rand(batch_size, device=image_tokens.device)
+        # Sample a random mask probability for each image using timestep and cosine schedule
+        mask_prob = mask_schedule(timesteps)
+        mask_prob = mask_prob.clip(config.training.min_masking_rate)
+
+    # creat a random mask for each image
+    num_token_masked = (seq_len * mask_prob).round().clamp(min=1)
+
+    mask_contiguous_region_prob = config.training.get("mask_contiguous_region_prob", None)
+
+    if mask_contiguous_region_prob is None:
+        mask_contiguous_region = False
+    else:
+        mask_contiguous_region = random.random() < mask_contiguous_region_prob
+
+    if not mask_contiguous_region:
+        batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)
+        mask = batch_randperm < num_token_masked.unsqueeze(-1)
+    else:
+        resolution = int(seq_len ** 0.5)
+        mask = torch.zeros((batch_size, resolution, resolution), device=image_tokens.device)
+
+        # TODO - would be nice to vectorize
+        for batch_idx, num_token_masked_ in enumerate(num_token_masked):
+            num_token_masked_ = int(num_token_masked_.item())
+
+            # NOTE: a bit handwavy with the bounds but gets a rectangle of ~num_token_masked_
+            num_token_masked_height = random.randint(
+                math.ceil(num_token_masked_ / resolution), min(resolution, num_token_masked_)
+            )
+            num_token_masked_height = min(num_token_masked_height, resolution)
+
+            num_token_masked_width = math.ceil(num_token_masked_ / num_token_masked_height)
+            num_token_masked_width = min(num_token_masked_width, resolution)
+
+            start_idx_height = random.randint(0, resolution - num_token_masked_height)
+            start_idx_width = random.randint(0, resolution - num_token_masked_width)
+
+            mask[
+            batch_idx,
+            start_idx_height: start_idx_height + num_token_masked_height,
+            start_idx_width: start_idx_width + num_token_masked_width,
+            ] = 1
+
+        mask = mask.reshape(batch_size, seq_len)
+        mask = mask.to(torch.bool)
+
+    # mask images and create input and labels
+    if config.training.get("noise_type", "mask"):
+        input_ids = torch.where(mask, mask_id, image_tokens)
+    elif config.training.get("noise_type", "random_replace"):
+        # sample random tokens from the vocabulary
+        random_tokens = torch.randint_like(
+            image_tokens, low=0, high=config.model.codebook_size, device=image_tokens.device
+        )
+        input_ids = torch.where(mask, random_tokens, image_tokens)
+    else:
+        raise ValueError(f"noise_type {config.training.noise_type} not supported")
+
+    if (
+            config.training.get("predict_all_tokens", False)
+            or config.training.get("noise_type", "mask") == "random_replace"
+    ):
+        labels = image_tokens
+        loss_weight = get_loss_weight(mask_prob, mask.long())
+    else:
+        labels = torch.where(mask, image_tokens, -100)
+        loss_weight = None
+
+    return input_ids, labels, loss_weight, mask_prob
+
+
+##################################################
+#              misc
+##################################################
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+from torchvision import transforms
+def image_transform(image, resolution=256, normalize=True):
+    image = transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR)(image)
+    image = transforms.CenterCrop((resolution, resolution))(image)
+    image = transforms.ToTensor()(image)
+    if normalize:
+        image = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)(image)
+    return image

validation_prompts/showoprompts.txt

@@ -0,0 +1,24 @@
+Close-up view of a computer screen, with the screen displaying a webpage.
+A tranquil scene of a lotus pond with koi fish swimming gracefully in a peaceful Chinese ink painting.
+Paper artwork, layered paper, colorful Chinese dragon surrounded by clouds.
+Pixel art character riding a dragon through the clouds.
+A peaceful village nestled at the foot of towering mountains in a tranquil East Asian watercolor scene.
+A person with swirling patterns of teal paint on their face and a shimmering silver crescent moon placed above their eyebrow, symbolizing mystery and magic.
+A dynamic scene of a rally car race.
+The breathtaking view of Santorini, a renowned landmark in Greece. The white-washed buildings with blue domes overlook the deep blue waters of the Aegean Sea, creating a stunning contrast against the vibrant sunset.
+An abstract portrait of a pensive face, rendered in cool shades of blues, purples, and grays.
+A punk rock frog in a studded leather jacket shouting into a microphone while standing on a boulder.
+A rebellious squirrel in a studded denim vest, strumming an electric guitar with fervor in a forest clearing.
+a captivating watercolor portrait of a dog's head, rendered in a vibrant palette of colors.
+A captivating watercolor portrait of a cat's face, rendered in a soft palette of pastels.
+A captivating watercolor portrait of a rabbit's profile, rendered in gentle hues of pinks and browns.
+a white Lamborghini Gallardo Spyder is parked on a cobblestone street.
+the breathtaking beauty of Whitehaven Beach.
+The breathtaking view of Moraine Lake, a renowned landmark in Canada. The turquoise waters of the lake reflect the rugged peaks of the Valley of the Ten Peaks, creating a scene of unparalleled natural beauty.
+The breathtaking view of Mount Fuji, a renowned landmark in Japan. The iconic snow-capped peak rises majestically above the surrounding landscape, mirrored perfectly in the tranquil waters of Lake Kawaguchi.
+A bustling Asian market at night, with colorful lanterns, street food vendors, and a mix of traditional and modern architecture.
+A stunning coastal cliffside at sunset, with waves crashing against the rocks and the sky painted in shades of orange, pink, and purple.
+A tranquil island paradise, with a white sandy beach, crystal-clear water, and palm trees swaying in the gentle breeze.
+A carnival of dreams where carousel horses gallop into the sky and cotton candy clouds drift by.
+A floating market in the sky where clouds serve as stalls for trading dreams.
+Intricate paper-cut creation featuring a vibrant peacock perched among blooming flowers.