Update app.py

app.py

@@ -21,9 +21,7 @@ CHECKPOINTS = {
 
 # 全局变量
 HF_TOKEN = os.getenv("HF_TOKEN")
-current_vis = []
-current_bpe = []
-current_index = 0
+
 
 
 def load_model(check_type):
@@ -55,7 +53,7 @@ def load_model(check_type):
     return model.to(device), tokenizer, transform, device
 
 def process_image(model, tokenizer, transform, device, check_type, image, text):
-    global current_vis, current_bpe, current_index
+    
     src_size = image.size
     if 'TokenOCR' in check_type:
         images, target_ratio = dynamic_preprocess(image, min_num=1, max_num=12, 
@@ -109,21 +107,13 @@ def process_image(model, tokenizer, transform, device, check_type, image, text):
     # current_bpe[-1] = 'Input text'
     current_bpe[-1] = text
 
-    return image, current_vis[0], current_bpe[0]
+    return image, current_vis, current_bpe
 
 # 事件处理函数
 def update_index(change):
-    global current_vis, current_bpe, current_index
-    
-    # 限制索引范围
-    current_index = max(0, min(len(current_vis) - 1, current_index + change))
-    
-    # 处理按钮可见性
-    prev_visible = current_index > 0
-    next_visible = current_index < len(current_vis) - 1
     
-    return current_vis[current_index], format_bpe_display(current_bpe[current_index]), prev_visible, next_visible
-
+    aa=1
+    return 1
 
 def format_bpe_display(bpe):
     # 使用HTML标签来设置字体大小、颜色，加粗，并居中
@@ -131,8 +121,6 @@ def format_bpe_display(bpe):
 
 
 
-
-
 # Gradio界面
 with gr.Blocks(title="BPE Visualization Demo") as demo:
     gr.Markdown("## BPE Visualization Demo - TokenFD基座模型能力可视化")
@@ -175,36 +163,38 @@ with gr.Blocks(title="BPE Visualization Demo") as demo:
     # 事件处理
     @spaces.GPU
     def on_run_clicked(model_type, image, text):
-        global current_vis, current_bpe, current_index
-        current_index = 0  # 重新处理新图片时，索引归零
-        image, vis, bpe = process_image(*load_model(model_type), model_type, image, text)
-        
-        if len(current_vis) > 1:
-            prev_visible, next_visible = True, True
-        else:
-            prev_visible, next_visible = False, False
-        print("len_current_vis",len(current_vis))
-        print("len_current_bpe",len(current_bpe))
-        print("current_vis",current_vis)
-        print("current_bpe",current_bpe)
-        return image, vis, format_bpe_display(bpe), prev_visible, next_visible
-    
 
+        current_index = 0  # Reset index when new image is processed
+        image, current_vis, current_bpe = process_image(*load_model(model_type), model_type, image, text)
+        
+        bpe_text = format_bpe_display(current_bpe)
+        
+        print("current_vis",len(current_vis))
+        print("current_bpe",len(current_bpe))
+        
+        return image, current_vis[0], current_bpe[0]
 
     
+    run_btn.click(
+        on_run_clicked,
+        inputs=[model_type, image_input, text_input],
+        outputs=[orig_img, heatmap, bpe_display, prev_btn, next_btn],  # 直接更新 prev_btn 和 next_btn
+    )
+    
     prev_btn.click(
-        lambda: update_index(-1),
-        outputs=[heatmap, bpe_display, prev_btn, next_btn]
+        lambda: (*update_index(-1), current_index),
+        outputs=[heatmap, bpe_display]
     )
     
     next_btn.click(
-        lambda: update_index(1),
-        outputs=[heatmap, bpe_display, prev_btn, next_btn]
+        lambda: (*update_index(1), current_index),
+        outputs=[heatmap, bpe_display]
     )
 
 
 
 
 
+
 if __name__ == "__main__":
     demo.launch()