add app_t2i.py

app_t2i.py

@@ -0,0 +1,209 @@
+from PIL import Image
+import gradio as gr
+from huggingface_hub import hf_hub_download, snapshot_download
+import torch
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision('high')
+setattr(torch.nn.Linear, 'reset_parameters', lambda self: None)
+setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None)
+import os
+import time
+import argparse
+from tokenizer_image.vq_model import VQ_models
+from models.gpt import GPT_models
+from models.generate import generate
+from t5 import T5Embedder
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+device = "cuda"
+
+model2ckpt = {
+    "GPT-XL": ("vq_ds16_t2i.pt", "t2i_XL_stage2_512.pt", 512),
+}
+
+def load_model(args):
+    ckpt_folder = './'
+    t5_folder = os.path.join(ckpt_folder, "flan-t5-xl")
+    if not os.path.exists(t5_folder):
+        os.makedirs(t5_folder, exist_ok=True)
+    vq_ckpt, gpt_ckpt, image_size = model2ckpt[args.gpt_model]
+    hf_hub_download(repo_id="peizesun/llamagen_t2i", filename=vq_ckpt, local_dir=ckpt_folder)
+    hf_hub_download(repo_id="peizesun/llamagen_t2i", filename=gpt_ckpt, local_dir=ckpt_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="config.json", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="pytorch_model-00001-of-00002.bin", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="pytorch_model-00002-of-00002.bin", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="pytorch_model.bin.index.json", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="special_tokens_map.json", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="spiece.model", local_dir=t5_folder)
+    hf_hub_download(repo_id="google/flan-t5-xl", filename="tokenizer_config.json", local_dir=t5_folder)
+    # create and load model
+    vq_model = VQ_models[args.vq_model](
+        codebook_size=args.codebook_size,
+        codebook_embed_dim=args.codebook_embed_dim)
+    vq_model.to(device)
+    vq_model.eval()
+    checkpoint = torch.load(f"{ckpt_folder}{vq_ckpt}", map_location="cpu")
+    vq_model.load_state_dict(checkpoint["model"])
+    del checkpoint
+    print(f"image tokenizer is loaded")
+
+    # create and load gpt model
+    precision = {'none': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}[args.precision]
+    latent_size = image_size // args.downsample_size
+    gpt_model = GPT_models[args.gpt_model](
+        vocab_size=args.codebook_size,
+        block_size=latent_size ** 2,
+        num_classes=args.num_classes,
+        cls_token_num=args.cls_token_num,
+        model_type=args.gpt_type,
+    ).to(device=device, dtype=precision)
+    
+    checkpoint = torch.load(f"{ckpt_folder}{gpt_ckpt}", map_location="cpu")
+    if args.from_fsdp: # fspd
+        model_weight = checkpoint
+    elif "model" in checkpoint:  # ddp
+        model_weight = checkpoint["model"]
+    elif "module" in checkpoint: # deepspeed
+        model_weight = checkpoint["module"]
+    elif "state_dict" in checkpoint:
+        model_weight = checkpoint["state_dict"]
+    else:
+        raise Exception("please check model weight")
+    # if 'freqs_cis' in model_weight:
+    #     model_weight.pop('freqs_cis')
+    gpt_model.load_state_dict(model_weight, strict=False)
+    gpt_model.eval()
+    del checkpoint
+    print(f"gpt model is loaded")
+
+    if args.compile:
+        print(f"compiling the model...")
+        gpt_model = torch.compile(
+            gpt_model,
+            mode="reduce-overhead",
+            fullgraph=True
+        ) # requires PyTorch 2.0 (optional)
+    else:
+        print(f"no need to compile model in demo") 
+
+    t5_model = T5Embedder(
+        device=device, 
+        local_cache=True, 
+        cache_dir=ckpt_folder, 
+        dir_or_name="flan-t5-xl",
+        torch_dtype=precision,
+        model_max_length=args.t5_feature_max_len,
+    )
+
+    return t5_model, vq_model, gpt_model, image_size
+
+
+def infer(cfg_scale, top_k, top_p, temperature, prompt, seed):
+    prompts = [prompt for _ in range(4)]
+    caption_embs, emb_masks = t5_model.get_text_embeddings(prompts)
+
+    if not args.no_left_padding:
+        print(f"processing left-padding...")    
+        # a naive way to implement left-padding
+        new_emb_masks = torch.flip(emb_masks, dims=[-1])
+        new_caption_embs = []
+        for idx, (caption_emb, emb_mask) in enumerate(zip(caption_embs, emb_masks)):
+            valid_num = int(emb_mask.sum().item())
+            print(f'  prompt {idx} token len: {valid_num}')
+            new_caption_emb = torch.cat([caption_emb[valid_num:], caption_emb[:valid_num]])
+            new_caption_embs.append(new_caption_emb)
+        new_caption_embs = torch.stack(new_caption_embs)
+    else:
+        new_caption_embs, new_emb_masks = caption_embs, emb_masks
+    c_indices = new_caption_embs * new_emb_masks[:,:, None]
+    c_emb_masks = new_emb_masks
+    qzshape = [len(c_indices), args.codebook_embed_dim, latent_size, latent_size]
+
+    t1 = time.time()
+    torch.manual_seed(seed)
+    index_sample = generate(
+        gpt_model, c_indices, latent_size ** 2,
+        c_emb_masks,
+        cfg_scale=cfg_scale, cfg_interval=args.cfg_interval,
+        temperature=temperature, top_k=top_k,
+        top_p=top_p, sample_logits=True, 
+        )
+    sampling_time = time.time() - t1
+    print(f"gpt sampling takes about {sampling_time:.2f} seconds.")    
+
+    t2 = time.time()
+    samples = vq_model.decode_code(index_sample, qzshape) # output value is between [-1, 1]
+    decoder_time = time.time() - t2
+    print(f"decoder takes about {decoder_time:.2f} seconds.")
+    # Convert to PIL.Image format:
+    samples = samples.mul(127.5).add_(128.0).clamp_(0, 255).permute(0, 2, 3, 1).to("cpu", torch.uint8).numpy()
+    samples = [Image.fromarray(sample) for sample in samples]
+    return samples
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--t5-path", type=str, default='.')
+parser.add_argument("--t5-feature-max-len", type=int, default=120)
+parser.add_argument("--t5-feature-dim", type=int, default=2048)
+parser.add_argument("--no-left-padding", action='store_true', default=False)
+parser.add_argument("--gpt-model", type=str, choices=list(GPT_models.keys()), default="GPT-XL")
+parser.add_argument("--gpt-type", type=str, choices=['c2i', 't2i'], default="t2i", help="class-conditional or text-conditional")
+parser.add_argument("--from-fsdp", action='store_true')
+parser.add_argument("--cls-token-num", type=int, default=120, help="max token number of condition input")
+parser.add_argument("--precision", type=str, default='bf16', choices=["none", "fp16", "bf16"]) 
+parser.add_argument("--compile", action='store_true', default=False)
+parser.add_argument("--vq-model", type=str, choices=list(VQ_models.keys()), default="VQ-16")
+parser.add_argument("--codebook-size", type=int, default=16384, help="codebook size for vector quantization")
+parser.add_argument("--codebook-embed-dim", type=int, default=8, help="codebook dimension for vector quantization")
+parser.add_argument("--downsample-size", type=int, choices=[8, 16], default=16)
+parser.add_argument("--num-classes", type=int, default=1000)
+parser.add_argument("--cfg-scale", type=float, default=7.5)
+parser.add_argument("--cfg-interval", type=float, default=-1)
+parser.add_argument("--seed", type=int, default=0)
+parser.add_argument("--top-k", type=int, default=2000,help="top-k value to sample with")
+parser.add_argument("--temperature", type=float, default=1.0, help="temperature value to sample with")
+parser.add_argument("--top-p", type=float, default=1.0, help="top-p value to sample with")
+args = parser.parse_args()
+
+t5_model, vq_model, gpt_model, image_size = load_model(args)
+latent_size = image_size // args.downsample_size
+
+examples = [
+    "A fluffy golden retriever puppy with big, soulful eyes sits in a sunlit garden, surrounded by colorful flowers and butterflies fluttering around its wagging tail.",
+    "A steaming bowl of Pho, filled with translucent rice noodles and thin slices of savory beef, topped with a heaping of fresh bean sprouts, a wedge of lime on the side, and a sprinkle of chopped green onions and cilantro.",
+    "An ethereal black and white landscape, where a solitary, sinuous black tree stands stark against a stark white snowy backdrop. Its branches twist intricately towards the sky, casting dramatic shadows on the untouched snow below.",        
+]
+
+with gr.Blocks() as demo:
+    gr.Markdown("<h1 style='text-align: center'>Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation</h1>")
+
+    with gr.Tabs():
+        with gr.TabItem('Generate'):
+            with gr.Row():
+                with gr.Column():
+                    cfg_scale = gr.Slider(minimum=1, maximum=25, step=0.1, value=7.5, label='Classifier-free Guidance Scale')
+                    top_k = gr.Slider(minimum=1, maximum=16384, step=1, value=4000, label='Top-K')
+                    top_p = gr.Slider(minimum=0., maximum=1.0, step=0.1, value=1.0, label="Top-P")
+                    temperature = gr.Slider(minimum=0., maximum=1.0, step=0.1, value=1.0, label='Temperature')
+                    seed = gr.Slider(minimum=0, maximum=1000, step=1, value=0, label='Seed')
+                    with gr.Row():
+                        text_prompt = gr.Textbox(
+                            label="Enter your prompt",
+                            show_label=False,
+                            max_lines=1,
+                            placeholder="Enter your prompt",
+                        )
+                    button = gr.Button("Generate", variant="primary")
+                    gr.Examples(
+                        label="Examples (select one example, and click Generate button)",
+                        examples=examples,
+                        inputs=text_prompt,
+                        # outputs=[result],
+                        # fn=generate,
+                    )
+                with gr.Column():
+                    output = gr.Gallery(label='Generated Images', height=700)
+                    button.click(infer, inputs=[cfg_scale, top_k, top_p, temperature, text_prompt, seed], outputs=[output])
+    demo.queue()
+    demo.launch(debug=True, share=True)