update

app.py

@@ -225,9 +225,9 @@ class Handler:
         self.vqgan.to(self.device)
         self.vqgan.eval()
 
-        self.optvq = VQModelHF.from_pretrained("BorelTHU/optvq-16x16x4")
-        self.optvq.to(self.device)
-        self.optvq.eval()
+        # self.optvq = VQModelHF.from_pretrained("BorelTHU/optvq-16x16x4")
+        # self.optvq.to(self.device)
+        # self.optvq.eval()
 
         # self.vae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers")
         # self.vae.to(self.device)
@@ -236,9 +236,13 @@ class Handler:
         # self.revq = ReVQ.from_pretrained("AndyRaoTHU/revq-512T")
         # self.revq.to(self.device)
         # self.revq.eval()
-        self.sftok = SFTok.from_pretrained("AndyRaoTHU/SFTok-B")
-        self.sftok.to(self.device)
-        self.sftok.eval()
+        self.sftok_b = SFTok.from_pretrained("AndyRaoTHU/SFTok-B")
+        self.sftok_b.to(self.device)
+        self.sftok_b.eval()
+
+        self.sftok_l = SFTok.from_pretrained("AndyRaoTHU/SFTok-L")
+        self.sftok_l.to(self.device)
+        self.sftok_l.eval()
         print("Models loaded successfully!")
 
     def tensor_v2_to_image(self, tensor):
@@ -276,17 +280,20 @@ class Handler:
             # revq_rec = self.vae.decode(revq_rec).sample
             # sftok_rec = revq_rec
 
-            encoded_tokens = self.sftok.encode(img)[1]["min_encoding_indices"]
-            # encoded_tokens, _ = self.sftok.encode(img)
-            sftok_rec = self.sftok.decode_tokens(encoded_tokens)
+            encoded_tokens_b = self.sftok_b.encode(img)[1]["min_encoding_indices"]
+            sftok_rec_b = self.sftok_b.decode_tokens(encoded_tokens_b)
+
+            encoded_tokens_l = self.sftok_l.encode(img)[1]["min_encoding_indices"]
+            sftok_rec_l = self.sftok_l.decode_tokens(encoded_tokens_l)
         
         # tensor to PIL image
         img = self.tensor_to_image(img)
         basevq_rec = self.tensor_v2_to_image(basevq_rec)
         vqgan_rec = self.tensor_v2_to_image(vqgan_rec)
-        sftok_rec = self.tensor_to_image(sftok_rec)
+        sftok_rec_b = self.tensor_to_image(sftok_rec_b)
+        sftok_rec_l = self.tensor_to_image(sftok_rec_l)
 
-        return basevq_rec, vqgan_rec, sftok_rec
+        return basevq_rec, vqgan_rec, sftok_rec_b, sftok_rec_l
 
 if __name__ == "__main__":
     # create the model handler
@@ -301,48 +308,54 @@ if __name__ == "__main__":
         outputs=[
             gr.Image(label="BaseVQ Reconstruction", type="numpy"),
             gr.Image(label="VQGAN Reconstruction", type="numpy"),
-            gr.Image(label="SFTok Reconstruction", type="numpy"),
+            gr.Image(label="SFTok-B Reconstruction", type="numpy"),
+            gr.Image(label="SFTok-L Reconstruction", type="numpy"),
         ],
         title="Demo 1: Image Reconstruction",
         description="Upload an image to see how different VQ models (BaseVQ, VQGAN, SFTok) reconstruct it from latent codes."
     )
 
-    with gr.Blocks() as demo2:
-        gr.Markdown("## Demo 2: Codebook Reset Strategy Visualization")
-        gr.Markdown("Visualizes codebook and data movement at different training steps with or without codebook reset strategy.")
+    # with gr.Blocks() as demo2:
+    #     gr.Markdown("## Demo 2: Codebook Reset Strategy Visualization")
+    #     gr.Markdown("Visualizes codebook and data movement at different training steps with or without codebook reset strategy.")
+
+    #     with gr.Row():
+    #         num_data = gr.Slider(label="num_data", value=16, minimum=10, maximum=20, step=1)
+    #         num_code = gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1)
 
-        with gr.Row():
-            num_data = gr.Slider(label="num_data", value=16, minimum=10, maximum=20, step=1)
-            num_code = gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1)
+    #     submit_btn = gr.Button("Run Visualization")
 
-        submit_btn = gr.Button("Run Visualization")
+    #     with gr.Column():  # 垂直输出
+    #         out_without_reset = gr.Image(label="Without Reset")
+    #         out_with_reset = gr.Image(label="With Reset")
 
-        with gr.Column():  # 垂直输出
-            out_without_reset = gr.Image(label="Without Reset")
-            out_with_reset = gr.Image(label="With Reset")
+    #     submit_btn.click(fn=draw_reset_result, inputs=[num_data, num_code], outputs=[out_without_reset, out_with_reset])
 
-        submit_btn.click(fn=draw_reset_result, inputs=[num_data, num_code], outputs=[out_without_reset, out_with_reset])
 
+    # with gr.Blocks() as demo3:
+    #     gr.Markdown("## Demo 3: Channel Multi-Group Strategy Visualization")
+    #     gr.Markdown("Visualizes codebook and data movement at different training steps with or without multi-group strategy.")
 
-    with gr.Blocks() as demo3:
-        gr.Markdown("## Demo 3: Channel Multi-Group Strategy Visualization")
-        gr.Markdown("Visualizes codebook and data movement at different training steps with or without multi-group strategy.")
+    #     with gr.Row():
+    #         num_data = gr.Slider(label="num_data", value=32, minimum=28, maximum=40, step=1)
+    #         num_code = gr.Slider(label="num_code", value=8, minimum=6, maximum=10, step=1)
 
-        with gr.Row():
-            num_data = gr.Slider(label="num_data", value=32, minimum=28, maximum=40, step=1)
-            num_code = gr.Slider(label="num_code", value=8, minimum=6, maximum=10, step=1)
+    #     submit_btn = gr.Button("Run Visualization")
 
-        submit_btn = gr.Button("Run Visualization")
+    #     with gr.Column():  # 垂直输出
+    #         out_s = gr.Image(label="Single Group")
+    #         out_m = gr.Image(label="Multi Group")
 
-        with gr.Column():  # 垂直输出
-            out_s = gr.Image(label="Single Group")
-            out_m = gr.Image(label="Multi Group")
+    #     submit_btn.click(fn=draw_multi_group_result, inputs=[num_data, num_code], outputs=[out_s, out_m])
 
-        submit_btn.click(fn=draw_multi_group_result, inputs=[num_data, num_code], outputs=[out_s, out_m])
+    # demo = gr.TabbedInterface(
+    #     interface_list=[demo1, demo2, demo3],
+    #     tab_names=["Image Reconstruction", "Reset Strategy", "Channel Multi-Group Strategy"]
+    # )
 
     demo = gr.TabbedInterface(
-        interface_list=[demo1, demo2, demo3],
-        tab_names=["Image Reconstruction", "Reset Strategy", "Channel Multi-Group Strategy"]
+        interface_list=[demo1],
+        tab_names=["Image Reconstruction"]
     )
 
     demo.launch(share=True)