Inference code for text-image-to-image (TI→I) tasks

[fengxin-zhxx]

Hello, thank you for open-sourcing the Tar project.

I'm currently trying to reproduce the text-image-to-image (TI→I) task mentioned in your paper, aiming to generate images like "a photo of XX in the style of ".

My implementation is similar to this:

prompt = "a photo of panda in the style of <image>"
image = Image.open(image_path).convert('RGB')
image = to_tensor(image).unsqueeze(0).to(self.device)
image_code = self.input_visual_tokenizer(image)['encoded']
image_text = "".join([f"<I{x}>" for x in image_code[0].cpu().tolist()])

# Prepare prompt
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": f"{image_text}\n{prompt}"}
]

input_text = self.tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True)
input_text += f"<im_start><S{self.config.scale}>"
# Generate tokens
inputs = self.tokenizer(input_text, return_tensors="pt")
gen_ids = self.model.generate(
    inputs.input_ids.to(self.device),
    max_new_tokens=self.config.seq_len,
    do_sample=True,
    temperature=self.config.temperature,
    top_p=self.config.top_p,
    top_k=self.config.top_k)

# Process generated tokens
gen_text = self.tokenizer.batch_decode(gen_ids)[0]
gen_code = [int(x) for x in re.findall(r'<I(\d+)>', gen_text)]
gen_code = gen_code[:self.config.seq_len] + [0] * max(0, self.config.seq_len - len(gen_code))
gen_code = torch.tensor(gen_code).unsqueeze(0).to(self.device)

gen_tensor = self.visual_tokenizer.decode_from_encoder_indices(
    gen_code, 
    {'cfg_scale': self.config.cfg_scale}
)
gen_image = Image.fromarray(gen_tensor[0].numpy())

However, the generated image is almost identical to the input image, rather than a new image in the specified style. Could you please provide some guidance on the recommended inference code for handling subject-driven TI→I tasks with the Tar model?

Any code examples or detailed instructions would be greatly appreciated.

Thank you again for your hard work and for sharing this project with us!

[csuhan]

Hi, you can try to replace "a photo of panda in the style of " with "a photo of panda in the style of <I1> <I2> ..."

[fengxin-zhxx]

Thanks, but that doesn't seem to make a difference.. @csuhan

[csuhan]

Hi @fengxin-zhxx , we need some changes to the T2I script.

First, the gen_ids contains both the input and output tokens, so we need to strip them out.
gen_text = self.tokenizer.batch_decode(gen_ids[:, inputs.input_ids.shape[-1]:])[0]

Secondly, our released models are finetuned on T2I and I2T data, the model’s TI2I capability is being degraded. Here we provide a pretraining checkpoint for your reference (may not the base model of Tar-7B, I am not very sure, but they share similar pretraining data).

https://huggingface.co/csuhan/Tar-7B-Pretrain-temp

[csuhan]

Below is the TI2I scirpt.

import re
from dataclasses import dataclass

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms.functional import to_tensor
from transformers import AutoTokenizer, Qwen2ForCausalLM

from tok.mm_autoencoder import MMAutoEncoder


@dataclass
class TI2IConfig:
    model_path: str = "csuhan/Tar-7B-Pretrain-temp"
    # visual tokenizer config
    ar_path: str = 'ar_dtok_lp_256px.pth'
    encoder_path: str = 'ta_tok.pth'
    decoder_path: str = 'vq_ds16_t2i.pt'
    
    device: str = "cuda:0"
    dtype: torch.dtype = torch.bfloat16
    # generation parameters
    scale: int = 0  # choose from [0, 1, 2]
    seq_len: int = 729  # choose from [729, 169, 81]
    temperature: float = 1.0
    top_p: float = 0.95
    top_k: int = 1200
    cfg_scale: float = 4.0

class TextImageToImageInference:
    def __init__(self, config: TI2IConfig):
        self.config = config
        self.device = torch.device(config.device)
        self._load_models()
        
    def _load_models(self):
        self.model = Qwen2ForCausalLM.from_pretrained(self.config.model_path, torch_dtype=self.config.dtype).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)
        
        # Initialize visual tokenizer
        config = dict(
            ar_path=self.config.ar_path,
            encoder_path=self.config.encoder_path,
            decoder_path=self.config.decoder_path,
            encoder_args={'input_type': 'rec'},
            decoder_args={},
        )
        self.visual_tokenizer = MMAutoEncoder(**config).eval().to(dtype=self.config.dtype, device=self.device)
        self.visual_tokenizer.ar_model.cls_token_num = self.config.seq_len
        self.visual_tokenizer.encoder.pool_scale = self.config.scale + 1

    def generate_image(self, prompt: str, image_path: Image.Image) -> Image.Image:
        image = Image.open(image_path).convert('RGB')
        image = to_tensor(image).unsqueeze(0).to(self.device)
        image_code = self.visual_tokenizer.encoder(image.to(self.config.dtype))['bottleneck_rep']
        image_text = "".join([f"<I{x}>" for x in image_code[0].cpu().tolist()])
        # replace <image> with its real codes
        prompt = prompt.replace("<image>", image_text)
        # Prepare prompt
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
        
        input_text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True)
        input_text += f"<im_start><S{self.config.scale}>"
        
        # Generate tokens
        inputs = self.tokenizer(input_text, return_tensors="pt")
        gen_ids = self.model.generate(
            inputs.input_ids.to(self.device),
            max_new_tokens=self.config.seq_len,
            do_sample=True,
            temperature=self.config.temperature,
            top_p=self.config.top_p,
            top_k=self.config.top_k)
        
        # NOTE: strip the prefix
        gen_text = self.tokenizer.batch_decode(gen_ids[:, inputs.input_ids.shape[-1]:])[0]
        gen_code = [int(x) for x in re.findall(r'<I(\d+)>', gen_text)]
        gen_code = gen_code[:self.config.seq_len] + [0] * max(0, self.config.seq_len - len(gen_code))
        gen_code = torch.tensor(gen_code).unsqueeze(0).to(self.device)

        gen_tensor = self.visual_tokenizer.decode_from_encoder_indices(
            gen_code, 
            {'cfg_scale': self.config.cfg_scale}
        )
        gen_image = Image.fromarray(gen_tensor[0].numpy())
        return gen_image

def main():
    config = TI2IConfig()
    config.ar_path = hf_hub_download("csuhan/TA-Tok", "ar_dtok_lp_512px.pth")
    config.encoder_path = hf_hub_download("csuhan/TA-Tok", "ta_tok.pth")
    config.decoder_path = hf_hub_download("peizesun/llamagen_t2i", "vq_ds16_t2i.pt")
    inference = TextImageToImageInference(config)
    
    ref_image = "path/to/your/image.png"
    prompt = "A photo of a macaw in the style of <image>"
    image = inference.generate_image(prompt, ref_image)
    image.save("generated_image.png")

if __name__ == "__main__":
    main()

[fengxin-zhxx]

Hi @csuhan, I have got the results I wanted, thank you for your reply!