debug

app.py

@@ -125,6 +125,13 @@ def draw_process(x, y, std):
 
 
 # ReVQ: for reset strategy
+def fig_to_array(fig):
+    buf = BytesIO()
+    fig.savefig(buf, format='png')  # 改为 png，不用 webp
+    buf.seek(0)
+    image = Image.open(buf)
+    return np.array(image)
+
 def get_codebook(quantizer):
     with torch.no_grad():
         codes = quantizer.embeddings.squeeze().detach()
@@ -196,7 +203,10 @@ def draw_reset_result(num_data=16, num_code=12):
     fig_reset.suptitle("VQ Codebook Training with Reset", fontsize=24, y=1.05)
     fig_nreset.suptitle("VQ Codebook Training without Reset", fontsize=24, y=1.05)
 
-    return fig_reset, fig_nreset
+    img_reset = fig_to_array(fig_reset)
+    img_nreset = fig_to_array(fig_nreset)
+
+    return img_reset, img_nreset
 
 # end
 
@@ -294,8 +304,8 @@ if __name__ == "__main__":
             gr.Slider(label="num_code", value=12, minimum=8, maximum=16, step=1),
         ],
         outputs=[
-            gr.Plot(label="With Reset"),
-            gr.Plot(label="Without Reset")
+            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."