FoundationVision/LlamaGen

1from PIL import Image2import gradio as gr3from imagenet_en_cn import IMAGENET_1K_CLASSES4from huggingface_hub import hf_hub_download5import torch6torch.backends.cuda.matmul.allow_tf32 = True7torch.backends.cudnn.allow_tf32 = True8torch.set_float32_matmul_precision('high')9setattr(torch.nn.Linear, 'reset_parameters', lambda self: None)10setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None)11 12import time13import argparse14from tokenizer_image.vq_model import VQ_models15from models.gpt import GPT_models16from models.generate import generate17 18device = "cuda"19 20model2ckpt = {21    "GPT-XL": ("vq_ds16_c2i.pt", "c2i_XL_384.pt", 384),22    "GPT-B": ("vq_ds16_c2i.pt", "c2i_B_256.pt", 256),23}24 25def load_model(args):26    ckpt_folder = "./"27    vq_ckpt, gpt_ckpt, image_size = model2ckpt[args.gpt_model]28    hf_hub_download(repo_id="FoundationVision/LlamaGen", filename=vq_ckpt, local_dir=ckpt_folder)29    hf_hub_download(repo_id="FoundationVision/LlamaGen", filename=gpt_ckpt, local_dir=ckpt_folder)30    # create and load model31    vq_model = VQ_models[args.vq_model](32        codebook_size=args.codebook_size,33        codebook_embed_dim=args.codebook_embed_dim)34    vq_model.to(device)35    vq_model.eval()36    checkpoint = torch.load(f"{ckpt_folder}{vq_ckpt}", map_location="cpu")37    vq_model.load_state_dict(checkpoint["model"])38    del checkpoint39    print(f"image tokenizer is loaded")40 41    # create and load gpt model42    precision = {'none': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}[args.precision]43    latent_size = image_size // args.downsample_size44    gpt_model = GPT_models[args.gpt_model](45        vocab_size=args.codebook_size,46        block_size=latent_size ** 2,47        num_classes=args.num_classes,48        cls_token_num=args.cls_token_num,49        model_type=args.gpt_type,50    ).to(device=device, dtype=precision)51    52    checkpoint = torch.load(f"{ckpt_folder}{gpt_ckpt}", map_location="cpu")53    if args.from_fsdp: # fspd54        model_weight = checkpoint55    elif "model" in checkpoint:  # ddp56        model_weight = checkpoint["model"]57    elif "module" in checkpoint: # deepspeed58        model_weight = checkpoint["module"]59    elif "state_dict" in checkpoint:60        model_weight = checkpoint["state_dict"]61    else:62        raise Exception("please check model weight")63    # if 'freqs_cis' in model_weight:64    #     model_weight.pop('freqs_cis')65    gpt_model.load_state_dict(model_weight, strict=False)66    gpt_model.eval()67    del checkpoint68    print(f"gpt model is loaded")69 70    if args.compile:71        print(f"compiling the model...")72        gpt_model = torch.compile(73            gpt_model,74            mode="reduce-overhead",75            fullgraph=True76        ) # requires PyTorch 2.0 (optional)77    else:78        print(f"no need to compile model in demo") 79 80    return vq_model, gpt_model, image_size81 82 83def infer(cfg_scale, top_k, top_p, temperature, class_label, seed):84    n = 485    latent_size = image_size // args.downsample_size86    # Labels to condition the model with (feel free to change):87    class_labels = [class_label for _ in range(n)]88    c_indices = torch.tensor(class_labels, device=device)89    qzshape = [len(class_labels), args.codebook_embed_dim, latent_size, latent_size]90 91    t1 = time.time()92    torch.manual_seed(seed)93    index_sample = generate(94        gpt_model, c_indices, latent_size ** 2,95        cfg_scale=cfg_scale, cfg_interval=args.cfg_interval,96        temperature=temperature, top_k=top_k,97        top_p=top_p, sample_logits=True, 98        )99    sampling_time = time.time() - t1100    print(f"gpt sampling takes about {sampling_time:.2f} seconds.")    101 102    t2 = time.time()103    samples = vq_model.decode_code(index_sample, qzshape) # output value is between [-1, 1]104    decoder_time = time.time() - t2105    print(f"decoder takes about {decoder_time:.2f} seconds.")106    # Convert to PIL.Image format:107    samples = samples.mul(127.5).add_(128.0).clamp_(0, 255).permute(0, 2, 3, 1).to("cpu", torch.uint8).numpy()108    samples = [Image.fromarray(sample) for sample in samples]109    return samples110 111 112parser = argparse.ArgumentParser()113parser.add_argument("--gpt-model", type=str, choices=list(GPT_models.keys()), default="GPT-XL")114parser.add_argument("--gpt-type", type=str, choices=['c2i', 't2i'], default="c2i", help="class-conditional or text-conditional")115parser.add_argument("--from-fsdp", action='store_true')116parser.add_argument("--cls-token-num", type=int, default=1, help="max token number of condition input")117parser.add_argument("--precision", type=str, default='bf16', choices=["none", "fp16", "bf16"]) 118parser.add_argument("--compile", action='store_true', default=False)119parser.add_argument("--vq-model", type=str, choices=list(VQ_models.keys()), default="VQ-16")120parser.add_argument("--codebook-size", type=int, default=16384, help="codebook size for vector quantization")121parser.add_argument("--codebook-embed-dim", type=int, default=8, help="codebook dimension for vector quantization")122parser.add_argument("--downsample-size", type=int, choices=[8, 16], default=16)123parser.add_argument("--num-classes", type=int, default=1000)124parser.add_argument("--cfg-scale", type=float, default=4.0)125parser.add_argument("--cfg-interval", type=float, default=-1)126parser.add_argument("--seed", type=int, default=0)127parser.add_argument("--top-k", type=int, default=2000,help="top-k value to sample with")128parser.add_argument("--temperature", type=float, default=1.0, help="temperature value to sample with")129parser.add_argument("--top-p", type=float, default=1.0, help="top-p value to sample with")130args = parser.parse_args()131 132vq_model, gpt_model, image_size = load_model(args)133 134with gr.Blocks() as demo:135    gr.Markdown("<h1 style='text-align: center'>Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation</h1>")136 137    with gr.Tabs():138        with gr.TabItem('Generate'):139            with gr.Row():140                with gr.Column():141                    # with gr.Row():142                    #     image_size = gr.Radio(choices=[384], value=384, label='Peize Model Resolution')143                    with gr.Row():144                        i1k_class = gr.Dropdown(145                            list(IMAGENET_1K_CLASSES.values()),146                            value='Eskimo dog, husky [爱斯基摩犬,哈士奇]',147                            type="index", label='ImageNet-1K Class'148                        )149                    cfg_scale = gr.Slider(minimum=1, maximum=25, step=0.1, value=4.0, label='Classifier-free Guidance Scale')150                    top_k = gr.Slider(minimum=1, maximum=16384, step=1, value=4000, label='Top-K')151                    top_p = gr.Slider(minimum=0., maximum=1.0, step=0.1, value=1.0, label="Top-P")152                    temperature = gr.Slider(minimum=0., maximum=1.0, step=0.1, value=1.0, label='Temperature')153                    seed = gr.Slider(minimum=0, maximum=1000, step=1, value=42, label='Seed')154                    # seed = gr.Number(value=0, label='Seed')155                    button = gr.Button("Generate", variant="primary")156                with gr.Column():157                    output = gr.Gallery(label='Generated Images', height=700)158                    button.click(infer, inputs=[cfg_scale, top_k, top_p, temperature, i1k_class, seed], outputs=[output])159    demo.queue()160    demo.launch(debug=True)161