test

app.py

@@ -280,41 +280,60 @@ class Handler:
 
 if __name__ == "__main__":
     # create the model handler
-    handler = Handler(device=device)
-
-    print("Creating Gradio interface...")
-
-    # Demo 1 接口：图像重建
-    demo1 = gr.Interface(
-        fn=handler.process_image,
-        inputs=gr.Image(label="Input Image", type="numpy"),
-        outputs=[
-            gr.Image(label="BaseVQ Reconstruction", type="numpy"),
-            gr.Image(label="VQGAN Reconstruction", type="numpy"),
-            gr.Image(label="ReVQ Reconstruction", type="numpy"),
-        ],
-        title="Demo 1: Image Reconstruction",
-        description="Upload an image to see how different VQ models (BaseVQ, VQGAN, ReVQ) reconstruct it from latent codes."
-    )
-
-    demo2 = gr.Interface(
-        fn=draw_reset_result,  # 不再需要 handler 包装
-        inputs=[
-            gr.Slider(label="num_data", value=16, minimum=10, maximum=20, step=1),
-            gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1),
-        ],
-        outputs=[
-            gr.Image(label="With Reset"),
-            gr.Image(label="Without Reset")
-        ],
-        title="Demo 2: Codebook Reset Strategy Visualization",
-        description="Visualizes codebook and data movement at different training steps."
-    )
+    # handler = Handler(device=device)
+
+    # print("Creating Gradio interface...")
+
+    # # Demo 1 接口：图像重建
+    # demo1 = gr.Interface(
+    #     fn=handler.process_image,
+    #     inputs=gr.Image(label="Input Image", type="numpy"),
+    #     outputs=[
+    #         gr.Image(label="BaseVQ Reconstruction", type="numpy"),
+    #         gr.Image(label="VQGAN Reconstruction", type="numpy"),
+    #         gr.Image(label="ReVQ Reconstruction", type="numpy"),
+    #     ],
+    #     title="Demo 1: Image Reconstruction",
+    #     description="Upload an image to see how different VQ models (BaseVQ, VQGAN, ReVQ) reconstruct it from latent codes."
+    # )
+
+    # demo2 = gr.Interface(
+    #     fn=draw_reset_result,  # 不再需要 handler 包装
+    #     inputs=[
+    #         gr.Slider(label="num_data", value=16, minimum=10, maximum=20, step=1),
+    #         gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1),
+    #     ],
+    #     outputs=[
+    #         gr.Image(label="With Reset"),
+    #         gr.Image(label="Without Reset")
+    #     ],
+    #     title="Demo 2: Codebook Reset Strategy Visualization",
+    #     description="Visualizes codebook and data movement at different training steps."
+    # )
+    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 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")
+
+        with gr.Column():  # 垂直输出
+            out_with_reset = gr.Image(label="With Reset")
+            out_without_reset = gr.Image(label="Without Reset")
+
+        submit_btn.click(fn=draw_reset_result, inputs=[num_data, num_code], outputs=[out_with_reset, out_without_reset])
 
     # 合并两个 interface 成 Tabbed UI
+    # demo = gr.TabbedInterface(
+    #     interface_list=[demo1, demo2],
+    #     tab_names=["Image Reconstruction", "Reset Strategy"]
+    # )
     demo = gr.TabbedInterface(
-        interface_list=[demo1, demo2],
-        tab_names=["Image Reconstruction", "Reset Strategy"]
+        interface_list=[demo2],
+        tab_names=["Reset Strategy"]
     )
 
     demo.launch(share=True)