add 512px AR

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 from torchvision.transforms.functional import to_tensor
 from huggingface_hub import hf_hub_download, snapshot_download, login
 
+from tok.ar_dtok.ar_model import ARModel
 from t2i_inference import T2IConfig, TextToImageInference
 
 def generate_text(self, image: str, prompt: str) -> str:
@@ -29,16 +30,16 @@ def generate_text(self, image: str, prompt: str) -> str:
 login(token=os.getenv('HF_TOKEN'))
 config = T2IConfig()
 config.model = snapshot_download("csuhan/Tar-7B-v0.1")
-config.ar_path = hf_hub_download("csuhan/TA-Tok", "ar_dtok_lp_1024px.pth")
+config.ar_path = {
+    "1024px": hf_hub_download("csuhan/TA-Tok", "ar_dtok_lp_1024px.pth"),
+    "512px": 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 = TextToImageInference(config)
 
-def generate_image(prompt, top_p, top_k, cfg_scale):
-    config.top_p = top_p
-    config.top_k = top_k
-    config.cfg_scale = cfg_scale
-    image = inference.generate_image(prompt)
+def generate_image(prompt, resolution, top_p, top_k, cfg_scale):
+    image = inference.generate_image(prompt, resolution, top_p, top_k, cfg_scale)
     return image
 
 def clear_inputs_t2i():
@@ -68,6 +69,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
           with gr.Column(scale=1):
               prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt")
               with gr.Accordion("Advanced Settings", open=False):
+                resolution = gr.Choice(
+                    ["512px", "1024px"], value="1024px", label="Resolution"
+                )
                 top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
                 top_k = gr.Slider(1, 2000, value=1200, step=10, label="Top-k")
                 cfg_scale = gr.Slider(1.0, 20.0, value=4.0, step=0.5, label="CFG Scale")
@@ -79,7 +83,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
       generate_btn.click(
           generate_image, 
-          inputs=[prompt, top_p, top_k, cfg_scale], 
+          inputs=[prompt, resolution, top_p, top_k, cfg_scale], 
           outputs=output_image
       )
       clear_btn.click(

t2i_inference.py

@@ -13,7 +13,7 @@ from tok.mm_autoencoder import MMAutoEncoder
 class T2IConfig:
     model_path: str = "csuhan/Tar-1.5B"
     # visual tokenizer config
-    ar_path: str = 'ar_dtok_lp_256px.pth'
+    ar_path = None
     encoder_path: str = 'ta_tok.pth'
     decoder_path: str = 'vq_ds16_t2i.pt'
     
@@ -39,17 +39,18 @@ class TextToImageInference:
         
         # Initialize visual tokenizer
         config = dict(
-            ar_path=self.config.ar_path,
+            ar_path_dict=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
+        for ar_model in self.visual_tokenizer.ar_model.values():
+            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.Image:
+    def generate_image(self, prompt, resolution, top_p, top_k, cfg_scale) -> Image.Image:
         # Prepare prompt
         messages = [
             {"role": "system", "content": "You are a helpful assistant."},
@@ -69,8 +70,8 @@ class TextToImageInference:
             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)
+            top_p=top_p,
+            top_k=top_k)
         
         # Process generated tokens
         gen_text = self.tokenizer.batch_decode(gen_ids)[0]
@@ -80,21 +81,7 @@ class TextToImageInference:
         
         gen_tensor = self.visual_tokenizer.decode_from_encoder_indices(
             gen_code, 
-            {'cfg_scale': self.config.cfg_scale}
+            {'cfg_scale': cfg_scale, 'resolution': resolution},
         )
         gen_image = Image.fromarray(gen_tensor[0].numpy())
-        return gen_image
-
-def main():
-    config = T2IConfig()
-    config.ar_path = hf_hub_download("csuhan/TA-Tok", "ar_dtok_lp_1024px.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 = TextToImageInference(config)
-    
-    prompt = "A photo of a macaw"
-    image = inference.generate_image(prompt)
-    image.save("generated_image.png")
-
-if __name__ == "__main__":
-    main()
+        return gen_image

tok/mm_autoencoder.py

@@ -8,17 +8,18 @@ from tok.ta_tok import TextAlignedTokenizer
 
 class MMAutoEncoder(nn.Module):
     def __init__(self,
-        ar_path,
+        ar_path_dict,
         encoder_path, decoder_path, 
         encoder_args={}, decoder_args={}):
         super().__init__()
-        self.ar_model = ARModel.from_checkpoint(ar_path)
+        self.ar_model = {resolution: ARModel.from_checkpoint(ar_path) for resolution, ar_path in ar_path_dict.items()}
 
         self.encoder = TextAlignedTokenizer.from_checkpoint(encoder_path, load_teacher=False, **encoder_args)
         self.decoder = VQVAE.from_checkpoint(decoder_path, **decoder_args)
 
     def ar_sample(self, x, args):
-        x = self.ar_model.sample(
+        resolution = args.get("resolution", "1024px")
+        x = self.ar_model[resolution].sample(
             x,
             cfg_scale=args.get('cfg_scale', 1.0),
             cfg_interval=args.get('cfg_interval', -1),