enable app

app.py

@@ -1,64 +1,314 @@
+import time
+from threading import Thread
+
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+import PIL
+from PIL import Image
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+from tqdm import tqdm
+from helpers import sample, expand2square
+
+# 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()
+
 
+# 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")
+
+
+IMAGE_TOKEN_INDEX=-200
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+   <img src='file/Liquid_icon.png' style="width: 80%; max-width: 600px; height: auto; opacity: 0.5;"> 
+   <h1 style="font-size: 20px; margin-bottom: 1px; opacity: 0.55;">UniTok-MLLM-7B</h1>
+</div>
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+
+CSS ="""
+.contain { display: flex; flex-direction: column; }
+#component-0 { height: 100%; }
+#chatbot { flex-grow: 1; }
 """
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+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>
+</div>
+"""
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+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>
+"""
 
-    messages.append({"role": "user", "content": message})
+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>. """
 
-    response = ""
+ckpt = torch.load(r'D:\projects\liquid_app\UniTok\UniTok_weights\unitok_tokenizer\unitok_tokenizer.pth', 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()
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
 
-        response += token
-        yield response
+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
+def bot_streaming_I2T(message, history):
+    print(message)
+    global stop_flag
+    stop_flag = True
+    time.sleep(0.2)
+    stop_flag = False
+    torch.cuda.empty_cache()
+    if message["files"]:
+        # message["files"][-1] is a Dict or just a string
+        if type(message["files"][-1]) == dict:
+            image = message["files"][-1]["path"]
+        else:
+            image = message["files"][-1]
+    else:
+        # if there's no image uploaded for this turn, look for images in the past turns
+        # kept inside tuples, take the last one
+        for hist in history:
+            if type(hist[0]) == tuple:
+                image = hist[0][0]
+    try:
+        if image is None:
+            # Handle the case where image is None
+            gr.Error("You need to upload an image for LLaVA to work.")
+    except NameError:
+        # Handle the case where 'image' is not defined at all
+        gr.Error("You need to upload an image for LLaVA to work.")
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+    qs = message['text']
+    qs = '\x00<image>\x01' + '\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(image).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')
+        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
+            time.sleep(0.06)
+            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)
+    stop_flag = False
+    
+    text_inputs = [message]*4  # generate 4 samples once
+    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)
+
+    
+    with torch.no_grad():
+        if guidance_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)
+        
+        with torch.no_grad():
+            sampling_kwargs={'temperature': temperature, 'top_k': top_K, 'top_p': top_P, '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)   
+                outputs = vqllm.T2I_forward_withcache(
+                    **model_inputs,
+                    input_multi_ids=input_multi_ids,
+                    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[0])
+                    if guidance_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) * guidance_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)
+                    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)
+
+                pred_tokens.append(torch.cat(indices_arhead, dim=1))  # [numcodebook,bz*2]
+                input_multi_ids = torch.stack(pred_tokens, dim=-1)
+                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,
+                )
+            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]
+            
+        
+        generated_image_list = []
+        rec_images = vq_model.idx_to_img(image_vq_id)
+        for index, rec_image in enumerate(rec_images):
+            rec_img = PILtransform(rec_image.squeeze(0).add(1).mul_(0.5).clamp_(0, 1))
+            generated_image_list.append(rec_img)
+
+        torch.cuda.empty_cache()
+        yield show_gallery(generated_image_list)
+ 
+
+chatbot_T2I=gr.Chatbot(placeholder=PLACEHOLDER,height=600)
+chat_input_T2I = gr.Textbox(placeholder="Enter text prompts...", show_label=False)
+
+chatbot_I2T=gr.Chatbot(placeholder=PLACEHOLDER, scale=1)
+chat_input_I2T = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+
+
+with gr.Blocks(fill_height=True) as demo:
+
+    gr.Markdown(title_html)
+    gr.Markdown(links_html)  
+    gr.Markdown(introduction)
+
+    with gr.Tab("Text To Image"):
+
+        description="Enter a text prompt or simply try one of the examples below to generate 4 images at once. Click to display the full image. You can configure hyperparameters for image generation in the Advanced Settings. "
+        gr.Markdown(description)  
+        with gr.Accordion("⚙️ Advanced Settings", open=False):
+            with gr.Row():
+                guidance_scale = gr.Slider(1.0, 20.0, value=7.0, label="Guidance Scale")
+                temperature = gr.Slider(0.0, 1.0, value=0.9, label="temperature")
+                top_K = gr.Slider(1, 8192, value=4096, label="Top K")
+                top_P = gr.Slider(0.0, 1.0, value=0.99, label="Top P")
+ 
+        aaa = gr.ChatInterface(
+        fn=bot_streaming_T2I,
+        examples=[
+            ["young blue dragon with horn lightning in the style of dd fantasy full body",5.0, 0.9,4096,0.99],
+            ["A majestic Goddes of beauty, charming dressed in a regal, jeweled gown and ornate crown, her golden hair cascading down her back, in the style of Pino Daeni",5.0, 0.9,4096,0.99],
+            ["A highly realistic, closeup photograph of a beautiful 35 year old redread woman writing in her journal, sitting on her balcony wearing warm, stylish outfits. Shot on a Canon EOS R5, the image boasts sharp focus and intricate details. The heartwarming scene conveys love, connection, and the crisp winter atmosphere, dramatic lighting.",5.0, 0.9,4096,0.99],
+            ["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],
+                  ],
+        stop_btn="Stop Generation",
+        additional_inputs = [guidance_scale, temperature, top_K, top_P],
+        additional_inputs_accordion="⚙️ Advanced Settings",
+        multimodal=False,
+        textbox=chat_input_T2I,
+        chatbot=chatbot_T2I,
+        fill_height=True,
+        )
+ 
+    with gr.Tab("Image To Text"):
+        bbb = gr.ChatInterface(
+            fn=bot_streaming_I2T,
+            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 )

helpers.py

@@ -0,0 +1,75 @@
+import torch
+from torch.nn import functional as F
+from PIL import Image
+
+
+
+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 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
+    

inference_i2t.py

@@ -0,0 +1,110 @@
+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_new.py

@@ -0,0 +1,243 @@
+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)