Upload app.py

app.py

@@ -0,0 +1,97 @@
+# Reference: https://huggingface.co/spaces/FoundationVision/LlamaGen/blob/main/app.py
+from PIL import Image
+import gradio as gr
+from imagenet_classes import imagenet_idx2classname
+import torch
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+import time
+import demo_util
+from utils.train_utils import create_pretrained_tokenizer
+import os
+import spaces
+from huggingface_hub import hf_hub_download
+
+os.system("pip3 install -U numpy")
+
+
+hf_hub_download(repo_id="fun-research/TiTok", filename="maskgit-vqgan-imagenet-f16-256.bin", local_dir="./")
+hf_hub_download(repo_id="yucornetto/RAR", filename="rar_b.bin", local_dir="./")
+
+# @spaces.GPU
+def load_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # load config
+    rar_model_size = "rar_b"
+    config = demo_util.get_config("configs/training/generator/rar.yaml")
+    config.experiment.generator_checkpoint = f"{rar_model_size}.bin"
+    config.model.generator.hidden_size = {"rar_b": 768, "rar_l": 1024, "rar_xl": 1280, "rar_xxl": 1408}[rar_model_size]
+    config.model.generator.num_hidden_layers = {"rar_b": 24, "rar_l": 24, "rar_xl": 32, "rar_xxl": 40}[rar_model_size]
+    config.model.generator.num_attention_heads = 16
+    config.model.generator.intermediate_size = {"rar_b": 3072, "rar_l": 4096, "rar_xl": 5120, "rar_xxl": 6144}[rar_model_size]
+
+    print(config)
+    tokenizer = create_pretrained_tokenizer(config)
+    print(tokenizer)
+    generator = demo_util.get_rar_generator(config)
+    print(generator)
+
+    tokenizer = tokenizer.to(device)
+    generator = generator.to(device)
+    return tokenizer, generator
+
+tokenizer, generator = load_model()
+
+@spaces.GPU
+def demo_infer(
+               guidance_scale, randomize_temperature, guidance_scale_pow,
+               class_label, seed):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    n = 4
+    class_labels = [class_label for _ in range(n)]
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    t1 = time.time()
+    generated_image = demo_util.sample_fn(
+        generator=generator,
+        tokenizer=tokenizer,
+        labels=class_labels,
+        guidance_scale=guidance_scale,
+        randomize_temperature=randomize_temperature,
+        guidance_scale_pow=guidance_scale_pow,
+        device=device
+    )
+    sampling_time = time.time() - t1
+    print(f"generation takes about {sampling_time:.2f} seconds.")    
+    samples = [Image.fromarray(sample) for sample in generated_image]
+    return samples
+
+with gr.Blocks() as demo:
+    gr.Markdown("<h1 style='text-align: center'>An Image is Worth 32 Tokens for Reconstruction and Generation</h1>")
+
+    with gr.Tabs():
+        with gr.TabItem('Generate'):
+            with gr.Row():
+                with gr.Column():
+                    with gr.Row():
+                        i1k_class = gr.Dropdown(
+                            list(imagenet_idx2classname.values()),
+                            value='Eskimo dog, husky',
+                            type="index", label='ImageNet-1K Class'
+                        )
+                    guidance_scale = gr.Slider(minimum=1, maximum=25, step=0.1, value=10.0, label='Classifier-free Guidance Scale')
+                    randomize_temperature = gr.Slider(minimum=0.8, maximum=1.2, step=0.01, value=1.0, label='randomize_temperature')
+                    guidance_scale_pow = gr.Slider(minimum=0.0, maximum=4.0, step=0.25, value=0.0, label='guidance_scale_pow')
+                    seed = gr.Slider(minimum=0, maximum=1000, step=1, value=42, label='Seed')
+                    button = gr.Button("Generate", variant="primary")
+                with gr.Column():
+                    output = gr.Gallery(label='Generated Images',
+                                        columns=4,
+                                        rows=1,
+                                        height=256, object_fit="scale-down")
+                    button.click(demo_infer, inputs=[
+                        guidance_scale, randomize_temperature, guidance_scale_pow,
+                        i1k_class, seed],
+                        outputs=[output])
+    demo.queue()
+    demo.launch(debug=True)