update ui

app.py

@@ -1,33 +1,26 @@
 import time
-from threading import Thread
-
-import gradio as gr
 import torch
-import PIL
+import gradio as gr
 from PIL import Image
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
 from tqdm import tqdm
-from helpers import sample, expand2square
+from threading import Thread
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
-# from transformers import AutoProcessor, LlavaForConditionalGeneration
-from transformers import TextIteratorStreamer
-from conversation import  conv_templates
 from model import *
 from unitok.config import Args
 from unitok.model import UniTok
-from mm_utils import tokenizer_image_token, get_model_name_from_path
-from torchvision import transforms
-
-PILtransform = transforms.ToPILImage()
+from conversation import conv_templates
+from mm_utils import tokenizer_image_token
+from helpers import sample, expand2square
 
+import os
+os.system("pip uninstall -y gradio")
+os.system("pip install gradio==4.44.1")
+os.system("pip install gradio_client==1.3.0")
 
-# import spaces
-# import os 
-# os.system("pip uninstall -y gradio") 
-# os.system("pip install gradio==4.44.1")
-# os.system("pip install gradio_client==1.3.0")
+PILtransform = transforms.ToPILImage()
 
 
 IMAGE_TOKEN_INDEX=-200
@@ -47,18 +40,20 @@ CSS ="""
 
 title_html = """
 <div style="display: flex; flex-direction: column; align-items: center; gap: 10px;">
-<h1 style="margin: 0; line-height: 1; text-align: center;"> Liquid: Language Models are Scalable Multi-modal <br> Generators via Unified Understanding and Generation</h1>
+<h1 style="margin: 0; line-height: 1; text-align: center;">UniTok: A Unified Tokenizer for Visual Generation and Understanding</h1>
 </div>
 """
 
 links_html = f"""
-<center><font size=3><a href='https://foundationvision.github.io/Liquid/'>Liquid</a> has been open-sourced on <a href='https://huggingface.co/Junfeng5/Liquid_V1_7B'>😊 Huggingface</a> and <a href='https://github.com/FoundationVision/Liquid'>🌟 GitHub</a>. If you find Liquid useful, a like❤️ or a star🌟 would be appreciated.</font></center>
+<center><font size=3><a href='https://foundationvision.github.io/Liquid/'>UniTok</a> has been open-sourced on <a href='https://huggingface.co/FoundationVision/unitok_mllm'>😊 Huggingface</a> and <a href='https://github.com/FoundationVision/UniTok'>🌟 GitHub</a>. If you find Liquid useful, a like❤️ or a star🌟 would be appreciated.</font></center>
 """
 
 introduction = f"""
- Liquid explores the potential of a single LLM as a multimodal generator and its scaling laws. It achieves the level of diffusion models in visual generation and discovers the mutual enhancement between understanding and generation. More details can be found on the project <a href='https://foundationvision.github.io/Liquid/'> homepage</a> and in the <a href='https://arxiv.org/abs/2412.04332'> paper</a>. """
+This is a native MLLM built with UniTok, a unified visual tokenizer well-suited for both generation and understanding tasks.
+More details can be found on the project <a href='https://foundationvision.github.io/UniTok/'> homepage</a> and in the <a href='https://arxiv.org/abs/2502.20321'> paper</a>. """
 
-ckpt = torch.load(r'D:\projects\liquid_app\UniTok\UniTok_weights\unitok_tokenizer\unitok_tokenizer.pth', map_location='cpu')
+unitok_ckpt = hf_hub_download(repo_id='FoundationVision/unitok_tokenizer', filename='unitok_tokenizer.pth')
+ckpt = torch.load('unitok_tokenizer.pth', map_location='cpu')
 vae_cfg = Args()
 vae_cfg.load_state_dict(ckpt['args'])
 vq_model = UniTok(vae_cfg)
@@ -66,13 +61,12 @@ vq_model.load_state_dict(ckpt['trainer']['unitok'])
 vq_model.to('cuda')
 vq_model.eval()
 
+mllm_ckpt = 'FoundationVision/unitok_mllm'
+tokenizer = AutoTokenizer.from_pretrained(mllm_ckpt, padding_side='left')
+vqllm = MiniGeminiLlamaForCausalLM.from_pretrained(mllm_ckpt).cuda()
+vqllm = vqllm.to(dtype=torch.bfloat16)
+vqllm = vqllm.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(r'C:\debug_ckpts\unitok_mllm', padding_side='left')
-vqllm = MiniGeminiLlamaForCausalLM.from_pretrained(
-    r'C:\debug_ckpts\unitok_mllm',
-    attn_implementation='flash_attention_2',
-    torch_dtype=torch.bfloat16
-).to('cuda')
 num_codebooks = vae_cfg.num_codebooks
 
 # @spaces.GPU
@@ -149,14 +143,12 @@ def bot_streaming_I2T(message, history):
             yield generated_text
         
 
-
 def show_gallery(images):
     gallery = gr.Gallery(images, label="Gallery", columns=4, height="auto",preview=True,scale=0.05)  # 设置两行两列的布局
     return gallery
 
 # @spaces.GPU
 def bot_streaming_T2I(message, history,guidance_scale, temperature, top_K, top_P):
- 
     global stop_flag
     stop_flag = True
     time.sleep(0.2)
@@ -175,9 +167,7 @@ def bot_streaming_T2I(message, history,guidance_scale, temperature, top_K, top_P
             model_inputs = tokenizer(text_inputs + uncondition_text_inputs, return_tensors="pt", padding=True).to('cuda')
         else:
             model_inputs = tokenizer(text_inputs, return_tensors="pt", padding=True).to('cuda')
-        model_kwargs = {'attention_mask':model_inputs.pop('attention_mask'),
-                        'use_cache': True
-                        }
+        model_kwargs = {'attention_mask':model_inputs.pop('attention_mask'), 'use_cache': True}
         input_ids = model_inputs.pop('input_ids')
         batch_size, cur_len = input_ids.shape
         if "inputs_embeds" in model_kwargs:
@@ -239,7 +229,6 @@ def bot_streaming_T2I(message, history,guidance_scale, temperature, top_K, top_P
             del model_kwargs
             # image_vq_id = input_ids[:,prompt_length:prompt_length+256]-ori_vocabe_size
             image_vq_id = torch.stack(pred_tokens, dim=-1)[:ori_batchsize]
-            
         
         generated_image_list = []
         rec_images = vq_model.idx_to_img(image_vq_id)
@@ -284,10 +273,10 @@ with gr.Blocks(fill_height=True) as demo:
             ["Portrait of an asian woman. She has pink violet hair style with modern complex hairdressing. The background is dark with cyberpunk neon lights. Inspired by Cyberpunk 2077 and Blade Runner. Ultra realistic picture. To capture the image, you will use a fullframe DSLR or mirrorless camera with a highresolution sensor, an aperture of f2.8 or wider, and a shutter speed of 1500 second or faster. You will use natural light and reflectors to create a balanced and welllit image, and will experiment with different angles and compositions to create the most i",5.0, 0.9,4096,0.99],
             ["female character fantasy world, for fantasy story, protagonist, interesting and detailed clothes, beautiful, medieval fantasy  cinematic shot  photo taken by canon, photo taken by fuji, photo taken by kodak  incredibly detailed, sharpen, details  professional lighting , film lighting  350mm  lightroom  cinematography, hyper realism, cinematic, film quality",5.0, 0.9,4096,0.99],
             ["strawberries splashing, swirling liquid, realism, octane render, raytracing",5.0, 0.9,4096,0.99],
-                  ["hedgehog face, floating in space, wearing space suit no helmet, cinematic, 50mm f1.8, unreal engine 5",5.0, 0.9,4096,0.99],
-                  ["artificial intelligence, revolution, publishing, writer, hyperrealistic",5.0, 0.9,4096,0.99],
-                  ["A pig dressed as a mason, by Bill Gekas",5.0, 0.9,4096,0.99],
-                  ],
+            ["hedgehog face, floating in space, wearing space suit no helmet, cinematic, 50mm f1.8, unreal engine 5",5.0, 0.9,4096,0.99],
+            ["artificial intelligence, revolution, publishing, writer, hyperrealistic",5.0, 0.9,4096,0.99],
+            ["A pig dressed as a mason, by Bill Gekas",5.0, 0.9,4096,0.99],
+        ],
         stop_btn="Stop Generation",
         additional_inputs = [guidance_scale, temperature, top_K, top_P],
         additional_inputs_accordion="⚙️ Advanced Settings",
@@ -300,15 +289,13 @@ with gr.Blocks(fill_height=True) as demo:
     with gr.Tab("Image To Text"):
         bbb = gr.ChatInterface(
             fn=bot_streaming_I2T,
-            examples=[ {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
+            examples=[{"text": "How to make this pastry?", "files": ["./baklava.png"]}],
             description="Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
             stop_btn="Stop Generation",
             multimodal=True,
             textbox=chat_input_I2T,
             chatbot=chatbot_I2T,
             )
-        
-
 
 demo.queue(api_open=False)
 demo.launch(allowed_paths=["./"],server_port=2560, share=False )

i2t.py

@@ -1,192 +0,0 @@
-import os
-import sys
-import json
-import math
-import torch
-import argparse
-import shortuuid
-from tqdm import tqdm
-from PIL import Image
-from PIL import ImageFile
-from torchvision import transforms
-
-from constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from conversation import conv_templates, SeparatorStyle
-from model.builder import load_pretrained_model
-from tools import disable_torch_init
-from mm_utils import tokenizer_image_token, get_model_name_from_path
-from torch.utils.data import Dataset, DataLoader
-
-
-from unitok.config import Args
-from unitok.model import UniTok
-
-ImageFile.LOAD_TRUNCATED_IMAGES = False
-torch.set_grad_enabled(False)
-
-
-def split_list(lst, n):
-    """Split a list into n (roughly) equal-sized chunks"""
-    chunk_size = math.ceil(len(lst) / n)  # integer division
-    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]
-
-
-def get_chunk(lst, n, k):
-    chunks = split_list(lst, n)
-    return chunks[k]
-
-
-# Custom dataset class
-class CustomDataset(Dataset):
-    def __init__(self, questions, image_folder, tokenizer, image_processor, model_config):
-        self.questions = questions
-        self.image_folder = image_folder
-        self.tokenizer = tokenizer
-        self.image_processor = image_processor
-        self.model_config = model_config
-
-    def __getitem__(self, index):
-        line = self.questions[index]
-        image_file = line["image"]
-        qs = line["text"]
-
-        qs = DEFAULT_IMAGE_TOKEN + '\n' + qs
-
-        conv = conv_templates[args.conv_mode].copy()
-        conv.append_message(conv.roles[0], qs)
-        conv.append_message(conv.roles[1], None)
-        prompt = conv.get_prompt()
-        # prompt = prompt.replace('<image>','<boi><image><eoi>')
-        # import pdb;pdb.set_trace()
-        image = Image.open(os.path.join(self.image_folder, image_file)).convert('RGB')
-        # import pdb;pdb.set_trace()
-        pad_image = expand2square(image, (122, 116, 104) )
-        # import pdb;pdb.set_trace()
-        img = self.image_processor[0](pad_image).unsqueeze(0)
-        img = img.to('cuda')
-        # import pdb;pdb.set_trace()
-        with torch.no_grad():
-            vq_code = self.image_processor[1].img_to_idx(img)
-            vqcode = vq_code.cpu()
-
-        input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')
-
-
-        return input_ids,vqcode,os.path.join(self.image_folder, image_file) #, image_tensor, image_tensor_aux
-    
-    def __len__(self):
-        return len(self.questions)
-
-
-# DataLoader
-def create_data_loader(questions, image_folder, tokenizer, image_processor, model_config, batch_size=1, num_workers=0):
-    assert batch_size == 1, "batch_size must be 1"
-    dataset = CustomDataset(questions, image_folder, tokenizer, image_processor, model_config)
-    data_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
-    return data_loader
-
-def expand2square(pil_img, background_color):
-    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
-    
-def eval_model(args):
-    # Model
-    disable_torch_init()
-    model_path = os.path.expanduser(args.model_path)
-    model_name = get_model_name_from_path(model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, args.model_base, model_name, load_8bit=args.load_8bit)
-
-    ckpt = torch.load(args.tokenizer_path, map_location='cpu')
-    vae_cfg = Args()
-    vae_cfg.load_state_dict(ckpt['args'])
-    vq_model = UniTok(vae_cfg)
-    vq_model.load_state_dict(ckpt['trainer']['unitok'])
-    vq_model.to('cuda')
-    vq_model.eval()
-    del ckpt
-
-    crop_size = 256
-    transform = transforms.Compose([
-        transforms.Resize((crop_size, crop_size)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-    ])
-    image_processor = (transform, vq_model)
-
-    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
-    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
-    answers_file = os.path.expanduser(args.answers_file)
-    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
-    ans_file = open(answers_file, "w")
-
-    if 'plain' in args.conv_mode and 'finetune' not in model_name.lower() and 'mmtag' not in args.conv_mode:
-        args.conv_mode = args.conv_mode + '_mmtag'
-        print(f'It seems that this is a plain model, but it is not using a mmtag prompt, auto switching to {args.conv_mode}.')
-
-    data_loader = create_data_loader(questions, args.image_folder, tokenizer, image_processor, model.config)
-
-    for (input_ids, image_codes,imagepath), line in tqdm(zip(data_loader, questions), total=len(questions)):
-        idx = line["question_id"]
-        cur_prompt = line["text"]
-        
-        input_ids = input_ids.to(device=model.device, non_blocking=True)
-        image_codes = image_codes.to(device=model.device, non_blocking=True)
-        if hasattr(model, "update_prompt"):
-            model.update_prompt([[cur_prompt]])
-
-        with torch.inference_mode():
-            output_ids = model.generate_mllm(
-                input_ids,
-                images=image_codes,
-                images_aux=  None,
-                do_sample=True if args.temperature > 0 else False,
-                temperature=args.temperature,
-                top_p=args.top_p,
-                num_beams=args.num_beams,
-                max_new_tokens=args.max_new_tokens,
-                bos_token_id=tokenizer.bos_token_id,  # Begin of sequence token
-                eos_token_id=tokenizer.eos_token_id,  # End of sequence token
-                pad_token_id=tokenizer.pad_token_id,  # Pad token
-                use_cache=False
-            )
-
-        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
-        ans_id = shortuuid.uuid()
-        ans_file.write(json.dumps({
-            "question_id": idx,
-            "prompt": cur_prompt,
-            "text": outputs,
-            "answer_id": ans_id,
-            "model_id": model_name,
-            "metadata": {}
-        }) + "\n")
-    ans_file.close()
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
-    parser.add_argument("--tokenizer-path", type=str, required=True)
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--image-folder", type=str, default="")
-    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
-    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
-    parser.add_argument("--conv-mode", type=str, default="llava_v1")
-    parser.add_argument("--num-chunks", type=int, default=1)
-    parser.add_argument("--chunk-idx", type=int, default=0)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--top_p", type=float, default=None)
-    parser.add_argument("--num_beams", type=int, default=1)
-    parser.add_argument('--load_8bit', type=bool, default=False)
-    parser.add_argument("--max_new_tokens", type=int, default=128)
-    args = parser.parse_args()
-
-    eval_model(args)

inference_i2t.py

@@ -1,110 +0,0 @@
-import torch
-import argparse
-import PIL
-from PIL import Image
-import os
-from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from conversation import conv_templates, SeparatorStyle
-from torchvision import transforms
-
-from constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-
-from threading import Thread
-from unitok.config import Args
-from unitok.model import UniTok
-
-from model.builder import load_pretrained_model
-from mm_utils import tokenizer_image_token, get_model_name_from_path
-
-
-IMAGE_TOKEN_INDEX=-200
-
-def expand2square(pil_img, background_color):
-    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
-    
-
-def main(args):
-    
-    ckpt = torch.load(args.unitok_path, map_location='cpu')
-    vae_cfg = Args()
-    vae_cfg.load_state_dict(ckpt['args'])
-    vq_model = UniTok(vae_cfg)
-    vq_model.load_state_dict(ckpt['trainer']['unitok'])
-    vq_model.to('cuda')
-    vq_model.eval()
-
-    model_path = os.path.expanduser(args.mllm_path)
-    model_name = get_model_name_from_path(model_path)
-    tokenizer, vqllm, image_processor, context_len = load_pretrained_model(model_path, model_name, load_8bit=args.load_8bit)
-
-    qs = args.prompt
-    qs = '<boi><image><eoi>' + '\n' + qs
-    conv = conv_templates['llava_v1'].copy()
-    conv.append_message(conv.roles[0], qs)
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt()
-    
-    crop_size = 256
-    transform = transforms.Compose([
-        transforms.Resize((crop_size, crop_size)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-    ])
-    
-    print(prompt)
-    image = Image.open(args.image_path).convert('RGB')
-    pad_image = expand2square(image, (122, 116, 104) )
-    # import pdb;pdb.set_trace()
-    img = transform(pad_image).unsqueeze(0)
-    img = img.to('cuda')
-    # import pdb;pdb.set_trace()
-    with torch.no_grad():
-        vq_code = vq_model.img_to_idx(img)
-        image_codes = vq_code.unsqueeze(0)
-        
-        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt')
-        
-        # input_ids = torch.cat(text_ids, dim=0)
-        # input_embeddings = vqllm.embed_tokens(input_ids)
-        inputs =  {
-            "inputs":input_ids.unsqueeze(0).to("cuda:0"),
-            "images":image_codes.to("cuda:0"),
-            "max_new_tokens":1024,
-            "bos_token_id":tokenizer.bos_token_id,  # Begin of sequence token
-            "eos_token_id":tokenizer.eos_token_id,  # End of sequence token
-            "pad_token_id":tokenizer.pad_token_id,  # Pad token
-            }
-        streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True})
-
-        # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.
-        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-        thread = Thread(target=vqllm.generate_mllm, kwargs=generation_kwargs)
-        thread.start()
-        generated_text = ""
-        for new_text in streamer:
-            generated_text += new_text
-        print(generated_text)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Process some integers.')
-    parser.add_argument('--unitok_path', type=str, default=r'D:\projects\liquid_app\UniTok\UniTok_weights\unitok_tokenizer\unitok_tokenizer.pth',required=False)
-    parser.add_argument('--mllm_path', type=str, default= r'C:\debug_ckpts\unitok_mllm', required=False)
-    parser.add_argument('--prompt', type=str, required=True, help='input text prompt')
-    parser.add_argument('--image_path', type=str, required=True, help='input image path')
-    parser.add_argument('--load_8bit',  action='store_true', default=False, help='use 8bit to save memory')
-
-    args = parser.parse_args()
-    main(args)
-
-

t2i.py

@@ -1,224 +0,0 @@
-import os
-import sys
-import torch
-import argparse
-import numpy as np
-from tqdm import tqdm
-from torchvision import transforms
-from torch.nn import functional as F
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-from model import *
-from unitok.config import Args
-from unitok.model import UniTok
-
-
-PILtransform = transforms.ToPILImage()
-
-
-def top_k_top_p_filtering(
-    logits,
-    top_k: int = 0,
-    top_p: float = 1.0,
-    filter_value: float = -float("Inf"),
-    min_tokens_to_keep: int = 1,
-):
-    """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
-    # import pdb;pdb.set_trace()
-    return logits
-
-
-def sample(logits, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0, sample_logits=True):
-    logits = logits[:, -1, :] / max(temperature, 1e-5)
-    if top_k > 0 or top_p < 1.0:
-        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
-    probs = F.softmax(logits, dim=-1)
-    if sample_logits:
-        idx = torch.multinomial(probs, num_samples=1)
-    else:
-        _, idx = torch.topk(probs, k=1, dim=-1)
-    return idx, probs
-
-
-def split_list(input_list, chunk_size):
-    return [input_list[i:i + chunk_size] for i in range(0, len(input_list), chunk_size)]
-
-
-def get_args_parser():
-    parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
-    parser.add_argument('--unitok_path', type=str, required=True)
-    parser.add_argument('--mllm_path', type=str, required=True)
-    parser.add_argument('--prompt_file', type=str, required=True)
-    parser.add_argument('--result_dir', type=str, required=True)
-    parser.add_argument('--idx', type=int, default=0)
-    parser.add_argument('--tau', type=float, default=0.9)
-    parser.add_argument('--topk', type=int, default=2048)
-    parser.add_argument('--topp', type=float, default=0.96)
-    parser.add_argument('--cfg_scale', type=float, default=5.0)
-    return parser
-
-
-def main(args):
-    text_set_id = args.idx
-    tau = args.tau
-    topk = args.topk
-    topp = args.topp
-    cfg_scale = args.cfg_scale
-
-    print('loading vq model ...')
-    ckpt = torch.load(args.unitok_path, map_location='cpu')
-    vae_cfg = Args()
-    vae_cfg.load_state_dict(ckpt['args'])
-    vq_model = UniTok(vae_cfg)
-    vq_model.load_state_dict(ckpt['trainer']['unitok'])
-    vq_model.to('cuda')
-    vq_model.eval()
-
-    image_save_pth = '{}/GenAI-cfg_{}-topk_{}-topp_{}-tau_{}'.format(args.result_dir, str(cfg_scale), str(topk), str(topp), str(tau))
-
-    tokenizer = AutoTokenizer.from_pretrained(args.mllm_path, padding_side='left')
-    vqllm = AutoModelForCausalLM.from_pretrained(
-        args.mllm_path,
-        attn_implementation='flash_attention_2',
-        torch_dtype=torch.bfloat16
-    ).to('cuda')
-
-    num_processes = 8
-    chunk_size = 8  # batchsize
-    num_codebooks = vae_cfg.num_codebooks
-
-    with open(args.prompt_file, 'r') as f:
-        lines = f.readlines()
-    all_prompts = []
-    for index, line in enumerate(lines):
-        all_prompts.append({'Index': str(index + 1).zfill(5), 'Prompt': line.strip()})
-
-    chunked_filenames = np.array_split(all_prompts, num_processes)
-    subset = chunked_filenames[text_set_id].tolist()
-    chunk_inputs = split_list(subset, chunk_size)
-    for chunk in tqdm(chunk_inputs):
-        text_inputs = [v['Prompt'] for v in chunk]
-        uncondition_text_inputs = ['<unconditional>'] * len(text_inputs)
-        for i in range(len(text_inputs)):
-            text_inputs[i] = text_inputs[i] + ' Generate an image based on this description.'
-        ori_batchsize = len(text_inputs)
-
-        save_list = []
-        if cfg_scale > 1:
-            model_inputs = tokenizer(text_inputs + uncondition_text_inputs, return_tensors="pt", padding=True).to('cuda')
-            total_batchsize = len(text_inputs + uncondition_text_inputs)
-            model_inputs['input_ids'] = torch.cat([
-                model_inputs['input_ids'],
-                torch.empty(total_batchsize, 1).fill_(3).to(model_inputs['input_ids'])
-            ], dim=1)
-            model_inputs['attention_mask'] = torch.cat([
-                model_inputs['attention_mask'],
-                torch.empty(total_batchsize, 1).fill_(1).to(model_inputs['attention_mask'])
-            ], dim=1)
-        else:
-            model_inputs = tokenizer(text_inputs, return_tensors="pt", padding=True).to('cuda')
-            total_batchsize = len(text_inputs)
-            model_inputs['input_ids'] = torch.cat([
-                model_inputs['input_ids'],
-                torch.empty(total_batchsize, 1).fill_(3).to(model_inputs['input_ids'])
-            ], dim=1)
-            model_inputs['attention_mask'] = torch.cat([
-                model_inputs['attention_mask'],
-                torch.empty(total_batchsize, 1).fill_(1).to(model_inputs['attention_mask'])
-            ], dim=1)
-        with torch.no_grad():
-            sampling_kwargs = {'temperature': tau, 'top_k': topk, 'top_p': topp, 'sample_logits': True}
-            pred_tokens = []
-            input_multi_ids = None
-            for _ in range(256):
-                outputs = vqllm.T2I_forward_nocache(
-                    **model_inputs,
-                    input_multi_ids=input_multi_ids,
-                    use_cache=None,
-                    return_dict=True,
-                    output_attentions=False,
-                    output_hidden_states=False,
-                )
-                next_embed = outputs['last_hidden_state'][:, -1:, :]
-
-                indices_arhead = []
-                for i_head in range(num_codebooks):
-                    ar_next_embed = vqllm.ar_head(
-                        inputs_embeds=next_embed,
-                        use_cache=False,
-                        output_attentions=False,
-                        output_hidden_states=False,
-                        return_dict=False,
-                    )
-                    next_token_logits = vqllm.ar_head.linear_head(ar_next_embed)
-                    if cfg_scale > 1:
-                        cond_logits, uncond_logits = torch.split(next_token_logits, len(next_token_logits) // 2, dim=0)
-                        cfg_logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
-                        half_next_token, _ = sample(cfg_logits, **sampling_kwargs)
-                        next_token = torch.cat([half_next_token, half_next_token])  # [bz,1]
-                    else:
-                        next_token, next_prob = sample(next_token_logits, **sampling_kwargs)
-
-                    indices_arhead.append(next_token)
-                    if i_head < num_codebooks - 1:
-                        predicted_embed = vqllm.ar_head.codebooks[i_head](next_token)
-                        next_embed = torch.cat([next_embed, predicted_embed], dim=1)
-
-                # update generated ids, model inputs, and length for next step
-                pred_tokens.append(torch.cat(indices_arhead, dim=1))  # [numcodebook,bz*2]
-                input_multi_ids = torch.stack(pred_tokens, dim=-1)
-
-            del sampling_kwargs, model_inputs, outputs
-
-            image_vq_id = torch.stack(pred_tokens, dim=-1)[:ori_batchsize]
-            save_list.append(image_vq_id)
-
-        torch.cuda.empty_cache()
-
-        print('decoding images ...')
-        if not os.path.exists(image_save_pth):
-            os.makedirs(image_save_pth)
-        for datainfo, vq_code in zip(chunk, save_list[0]):
-            idx = datainfo['Index']
-            new_gen_ids = vq_code.unsqueeze(0).to('cuda')
-            rec_image = vq_model.idx_to_img(new_gen_ids)
-            rec_img = PILtransform(rec_image.squeeze(0).add(1).mul_(0.5).clamp_(0, 1))
-            rec_img.save('{}/{}.jpg'.format(image_save_pth, str(idx)))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser('genai inference script', parents=[get_args_parser()])
-    args = parser.parse_args()
-    main(args)

t2i_new.py

@@ -1,243 +0,0 @@
-import os
-import sys
-import torch
-import argparse
-import numpy as np
-from tqdm import tqdm
-from torchvision import transforms
-from torch.nn import functional as F
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-from model import *
-from unitok.config import Args
-from unitok.model import UniTok
-
-
-PILtransform = transforms.ToPILImage()
-
-
-def top_k_top_p_filtering(
-    logits,
-    top_k: int = 0,
-    top_p: float = 1.0,
-    filter_value: float = -float("Inf"),
-    min_tokens_to_keep: int = 1,
-):
-    """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
-    # import pdb;pdb.set_trace()
-    return logits
-
-
-def sample(logits, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0, sample_logits=True):
-    logits = logits[:, -1, :] / max(temperature, 1e-5)
-    if top_k > 0 or top_p < 1.0:
-        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
-    probs = F.softmax(logits, dim=-1)
-    if sample_logits:
-        idx = torch.multinomial(probs, num_samples=1)
-    else:
-        _, idx = torch.topk(probs, k=1, dim=-1)
-    return idx, probs
-
-
-def split_list(input_list, chunk_size):
-    return [input_list[i:i + chunk_size] for i in range(0, len(input_list), chunk_size)]
-
-
-def get_args_parser():
-    parser = argparse.ArgumentParser('Set transformer detector', add_help=False)
-    parser.add_argument('--unitok_path', type=str, default=r'D:\projects\liquid_app\UniTok\UniTok_weights\unitok_tokenizer\unitok_tokenizer.pth',required=False)
-    parser.add_argument('--mllm_path', type=str, default= r'C:\debug_ckpts\unitok_mllm', required=False)
-    # parser.add_argument('--prompt_file', type=str, required=True)
-    # parser.add_argument('--result_dir', type=str, required=True)
-    parser.add_argument('--idx', type=int, default=0)
-    parser.add_argument('--tau', type=float, default=0.9)
-    parser.add_argument('--topk', type=int, default=2048)
-    parser.add_argument('--topp', type=float, default=0.96)
-    parser.add_argument('--cfg_scale', type=float, default=5.0)
-    return parser
-
-
-def main(args):
-    text_set_id = args.idx
-    tau = args.tau
-    topk = args.topk
-    topp = args.topp
-    cfg_scale = args.cfg_scale
-
-    print('loading vq model ...')
-    ckpt = torch.load(args.unitok_path, map_location='cpu')
-    vae_cfg = Args()
-    vae_cfg.load_state_dict(ckpt['args'])
-    vq_model = UniTok(vae_cfg)
-    vq_model.load_state_dict(ckpt['trainer']['unitok'])
-    vq_model.to('cuda')
-    vq_model.eval()
-
- 
-    tokenizer = AutoTokenizer.from_pretrained(args.mllm_path, padding_side='left')
-    vqllm = MiniGeminiLlamaForCausalLM.from_pretrained(
-        args.mllm_path,
-        attn_implementation='flash_attention_2',
-        torch_dtype=torch.bfloat16
-    ).to('cuda')
-    num_codebooks = vae_cfg.num_codebooks
-    # import pdb;pdb.set_trace() 
-    chunk_inputs = [[{'Prompt':'a dog in grasee'},{'Prompt':'a dog in grasee'},{'Prompt':'a dog in grasee'},{'Prompt':'a dog in grasee'}]]
-    for chunk in tqdm(chunk_inputs):
-        text_inputs = [v['Prompt'] for v in chunk]
-        uncondition_text_inputs = ['<unconditional>\x00'] * len(text_inputs)
-        for i in range(len(text_inputs)):
-            text_inputs[i] = text_inputs[i] + ' Generate an image based on this description.\x00'
-        ori_batchsize = len(text_inputs)
-
-        save_list = []
-        if cfg_scale > 1:
-            model_inputs = tokenizer(text_inputs + uncondition_text_inputs, return_tensors="pt", padding=True).to('cuda')
-        else:
-            model_inputs = tokenizer(text_inputs, return_tensors="pt", padding=True).to('cuda')
-           
-        
-        model_kwargs = {'attention_mask':model_inputs.pop('attention_mask'),
-                        'use_cache': True
-                        }
-        input_ids = model_inputs.pop('input_ids')
-        batch_size, cur_len = input_ids.shape
-        if "inputs_embeds" in model_kwargs:
-            cur_len = model_kwargs["inputs_embeds"].shape[1]
-        model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
-        # import pdb;pdb.set_trace()
-        
-        
-        with torch.no_grad():
-            sampling_kwargs = {'temperature': tau, 'top_k': topk, 'top_p': topp, 'sample_logits': True}
-
-            pred_tokens = []
-            input_multi_ids = None
-            for i in tqdm(range(256)):
-                model_inputs = vqllm.prepare_inputs_for_generation(input_ids, **model_kwargs)  #model_inputs['input_ids'], model_inputs['position_ids'], model_inputs['cache_position']
-                # import pdb;pdb.set_trace()  #model_inputs['position_ids'] , model_inputs['cache_position']
-                outputs = vqllm.T2I_forward_withcache(
-                    **model_inputs,
-                    input_multi_ids=input_multi_ids,
-                    return_dict=True,
-                    output_attentions=False,
-                    output_hidden_states=False,
-                )
-
-                # import pdb;pdb.set_trace()
-                next_embed = outputs['last_hidden_state'][:, -1:, :]
-                # next_token_logits = outputs.logits[:, -1:, :]
-                indices_arhead = []
-
-                # for i_head in range(num_codebooks):
-                #     ar_next_embed =  vqllm.ar_head(inputs_embeds=next_embed,use_cache=False,output_attentions=False, output_hidden_states=False, return_dict=False,)
-                #     next_token_logits = vqllm.ar_head.linear_head(ar_next_embed[:,-1:,:])
-                #     # import pdb;pdb.set_trace()
-                #     image_probs = F.softmax(next_token_logits, dim=-1)
-                #     _, image_idx = torch.topk(image_probs, k=1, dim=-1) #[numcodebook,256,1]
-                #     next_token = image_idx[:,:,0]
-                #     indices_arhead.append(next_token)
-                #     # import pdb;pdb.set_trace()
-                #     # pred_tokens.append(next_token)
-                #     if i_head<num_codebooks-1:
-                #         predicted_embed =  vqllm.ar_head.codebooks[i_head](next_token)
-                #         next_embed = torch.cat([next_embed,predicted_embed],dim=1)
-
-                for i_head in range(num_codebooks):
-                    ar_next_embed = vqllm.ar_head(
-                        inputs_embeds=next_embed,
-                        use_cache=False,
-                        output_attentions=False,
-                        output_hidden_states=False,
-                        return_dict=False,
-                        )
-                    next_token_logits = vqllm.ar_head.linear_head(ar_next_embed[0])
-                    if cfg_scale > 1:
-                        cond_logits, uncond_logits = torch.split(next_token_logits, len(next_token_logits) // 2, dim=0)
-                        cfg_logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
-                        half_next_token, _ = sample(cfg_logits, **sampling_kwargs)
-                        # pred_tokens.append(half_next_token)
-                        next_token = torch.cat([half_next_token, half_next_token])  # [bz,1]
-                    else:
-                        next_token, next_prob = sample(next_token_logits, **sampling_kwargs)
-                        # pred_tokens.append(next_token)
-                    # import pdb;pdb.set_trace()
-                    indices_arhead.append(next_token)
-                    if i_head < num_codebooks - 1:
-                        predicted_embed = vqllm.ar_head.codebooks[i_head](next_token)
-                        next_embed = torch.cat([next_embed, predicted_embed], dim=1)
-
-                # update generated ids, model inputs, and length for next step
-                # import pdb;pdb.set_trace()
-                pred_tokens.append(torch.cat(indices_arhead, dim=1))  # [numcodebook,bz*2]
-                input_multi_ids = torch.stack(pred_tokens, dim=-1)
-                # import pdb;pdb.set_trace()
-                fake_id = torch.zeros_like(input_ids[:,:1])
-                input_ids = torch.cat([input_ids, fake_id], dim=-1) # add fake id for cache
-                
-                model_kwargs = vqllm._update_model_kwargs_for_generation(
-                    outputs,
-                    model_kwargs,
-                    is_encoder_decoder=vqllm.config.is_encoder_decoder,
-                )
-                pass
-            del sampling_kwargs
-            del model_inputs
-            del outputs
-            del model_kwargs
-            # image_vq_id = input_ids[:,prompt_length:prompt_length+256]-ori_vocabe_size
-            image_vq_id = torch.stack(pred_tokens, dim=-1)[:ori_batchsize]
-            # print(set(image_vq_id.tolist()))
-            save_list.append(image_vq_id)
-
-        torch.cuda.empty_cache()
-
-        print('decoding images ...')
-        image_save_pth = 'visualresults' 
-        if not os.path.exists(image_save_pth):
-            os.makedirs(image_save_pth)
-        for datainfo, vq_code in zip(chunk, save_list[0]):
-            new_gen_ids = vq_code.unsqueeze(0).to('cuda')
-            # import pdb;pdb.set_trace()
-            rec_image = vq_model.idx_to_img(new_gen_ids)
-            rec_img = PILtransform(rec_image.squeeze(0).add(1).mul_(0.5).clamp_(0, 1))
-            rec_img.save('{}/{}.jpg'.format(image_save_pth, 'test'))
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser('genai inference script', parents=[get_args_parser()])
-    args = parser.parse_args()
-    main(args)