Update app.py

app.py

@@ -21,10 +21,6 @@ CHECKPOINTS = {
 
 # 全局变量
 HF_TOKEN = os.getenv("HF_TOKEN")
-current_vis = []
-current_bpe = []
-current_index = 0
-
 
 def load_model(check_type):
     # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -54,129 +50,52 @@ 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, 
-#                                                   image_size=model.config.force_image_size,
-#                                                   use_thumbnail=model.config.use_thumbnail,
-#                                                   return_ratio=True)
-#         pixel_values = torch.stack([transform(img) for img in images]).to(device)
-#     else:
-#         pixel_values = torch.stack([transform(image)]).to(device)
-#         target_ratio = (1, 1)
-
-    # # 文本处理
-    # text += ' '
-    # input_ids = tokenizer(text)['input_ids'][1:]
-    # input_ids = torch.tensor(input_ids, device=device)
-    
-    # # 获取嵌入
-    # with torch.no_grad():
-    #     if 'R50' in check_type:
-    #         text_embeds = model.language_embedding(input_ids)
-    #     else:
-    #         text_embeds = model.tok_embeddings(input_ids)
-        
-    #     vit_embeds, size1 = model.forward_tokenocr(pixel_values.to(torch.bfloat16).to(device))
-    #     print("vit_embeds",vit_embeds)
-    #     print("vit_embeds,shape",vit_embeds.shape)
-    #     print("target_ratio",target_ratio)
-    #     print("check_type",check_type)
-    #     vit_embeds, size2 = post_process(vit_embeds, target_ratio, check_type)
-        
-    #     # 计算相似度
-    #     text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
-    #     vit_embeds = vit_embeds / vit_embeds.norm(dim=-1, keepdim=True)
-    #     similarity = text_embeds @ vit_embeds.T
-    #     resized_size = size1 if size1 is not None else size2
-
-    # # print(f"text_embeds shape: {text_embeds.shape}, numel: {text_embeds.numel()}") # text_embeds shape: torch.Size([4, 2048]), numel: 8192
-    # # print(f"vit_embeds shape: {vit_embeds.shape}, numel: {vit_embeds.numel()}") # vit_embeds shape: torch.Size([9728, 2048]), numel: 19922944
-    # # print(f"similarity shape: {similarity.shape}, numel: {similarity.numel()}")# similarity shape: torch.Size([4, 9728]), numel: 38912
-
-
-    # # 生成可视化
-    # attn_map = similarity.reshape(len(text_embeds), resized_size[0], resized_size[1])
-    # # attn_map = similarity.reshape(len(text_embeds), *target_ratio)
-    # all_bpe_strings = [tokenizer.decode(input_id) for input_id in input_ids]
-    # current_vis = generate_similiarity_map([image], attn_map, 
-    #                                        [tokenizer.decode([i]) for i in input_ids], 
-    #                                        [], target_ratio, src_size)
-    
-    # current_bpe = [tokenizer.decode([i]) for i in input_ids]
-    # # current_bpe[-1] = 'Input text'
-    # current_bpe[-1] = text
-    # print("current_vis",len(current_vis))
-    # print("current_bpe",len(current_bpe))
-    # return image, current_vis[0], current_bpe[0]
-
-def process_image(model, tokenizer, transform, device, check_type, image, text):
-    global current_vis, current_bpe, current_index
+def process_image(model, tokenizer, transform, device, check_type, image, text, state):
     src_size = image.size
-
-    # Convert PIL Image to Tensor and move to the appropriate device
     if 'TokenOCR' in check_type:
-        # If dynamic preprocessing is required, handle differently
         images, target_ratio = dynamic_preprocess(image, min_num=1, max_num=12, 
                                                   image_size=model.config.force_image_size,
                                                   use_thumbnail=model.config.use_thumbnail,
                                                   return_ratio=True)
-        pixel_values = torch.stack([transform(img).to(device) for img in images])
+        pixel_values = torch.stack([transform(img) for img in images]).to(device)
     else:
-        # Standard image processing for a single image
-        pixel_values = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
+        pixel_values = torch.stack([transform(image)]).to(device)
         target_ratio = (1, 1)
 
+    # 文本处理
     text += ' '
-    input_ids = tokenizer(text, return_tensors='pt').input_ids.to(device)  # Ensure tokens are on the same device
-
+    input_ids = tokenizer(text)['input_ids'][1:]
+    input_ids = torch.tensor(input_ids, device=device)
+    
+    # 获取嵌入
     with torch.no_grad():
         if 'R50' in check_type:
             text_embeds = model.language_embedding(input_ids)
         else:
             text_embeds = model.tok_embeddings(input_ids)
-
-        # vit_embeds, size1 = model.forward_tokenocr(pixel_values)
-        vit_embeds, size1 = model.forward_tokenocr(pixel_values.to(torch.bfloat16))
-
+        
+        vit_embeds, size1 = model.forward_tokenocr(pixel_values.to(torch.bfloat16).to(device))
         vit_embeds, size2 = post_process(vit_embeds, target_ratio, check_type)
-
+        
+        # 计算相似度
         text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
         vit_embeds = vit_embeds / vit_embeds.norm(dim=-1, keepdim=True)
         similarity = text_embeds @ vit_embeds.T
         resized_size = size1 if size1 is not None else size2
 
     attn_map = similarity.reshape(len(text_embeds), resized_size[0], resized_size[1])
-    current_vis = generate_similiarity_map([image], attn_map, 
-                                           [tokenizer.decode([i]) for i in input_ids.squeeze()], 
-                                           [], target_ratio, src_size)
-
-    current_bpe = [tokenizer.decode([i]) for i in input_ids.squeeze()]
-    current_bpe[-1] = text
-    return image, current_vis[0], current_bpe[0]
-
-
-# 事件处理函数
-def update_index(change):
-    global current_vis, current_bpe, current_index
-    current_index = max(0, min(len(current_vis) - 1, current_index + change))
-    return current_vis[current_index], format_bpe_display(current_bpe[current_index])
-
-def format_bpe_display(bpe):
-    # 使用HTML标签来设置字体大小、颜色，加粗，并居中
-    return f"<div style='text-align:center; font-size:20px;'><strong>Current BPE: <span style='color:red;'>{bpe}</span></strong></div>"
-
-def update_slider_index(x):
-    global current_vis, current_bpe, current_index
-    print(f"x: {x}, current_vis length: {len(current_vis)}, current_bpe length: {len(current_bpe)}")
-    if 0 <= x < len(current_vis) and 0 <= x < len(current_bpe):
-        return current_vis[x], format_bpe_display(current_bpe[x])
-    else:
-        return None, "索引超出范围"
-
+    all_bpe_strings = [tokenizer.decode(input_id) for input_id in input_ids]
+    vis = generate_similiarity_map([image], attn_map, 
+                                   [tokenizer.decode([i]) for i in input_ids], 
+                                   [], target_ratio, src_size)
+    
+    bpe = [tokenizer.decode([i]) for i in input_ids]
+    bpe[-1] = text
 
+    # Store results in state
+    state['current_vis'] = vis
+    state['current_bpe'] = bpe
+    return image, vis[0], bpe[0], len(vis) - 1
 
 # Gradio界面
 with gr.Blocks(title="BPE Visualization Demo") as demo:
@@ -218,58 +137,48 @@ with gr.Blocks(title="BPE Visualization Demo") as demo:
             
             bpe_display = gr.Markdown("Current BPE: ", visible=False)
 
-    # 事件处理
-    # @spaces.GPU
-    # def on_run_clicked(model_type, image, text):
-    #     global current_vis, current_bpe, current_index
-    #     current_index = 0  # Reset index when new image is processed
-    #     image, vis, bpe = process_image(*load_model(model_type), model_type, image, text)
-    #     # Update the slider range and set value to 0
-    #     slider_max_val = len(current_bpe) - 1
-    #     bpe_text = format_bpe_display(bpe)
-    #     print("current_vis",len(current_vis))
-    #     print("current_bpe",len(current_bpe))
-    #     return image, vis, bpe_text, slider_max_val
+    state = gr.State(current_vis=[], current_bpe=[], current_index=0)
 
     @spaces.GPU
-    def on_run_clicked(model_type, image, text):
-        global current_vis, current_bpe, current_index
-        current_index = 0
-        model, tokenizer, transform, device = load_model(model_type)
-        image, vis, bpe = process_image(model, tokenizer, transform, device, model_type, image, text)
-        slider_max_val = len(current_bpe) - 1
+    def on_run_clicked(model_type, image, text, state):
+        image, vis, bpe, slider_max_val = process_image(*load_model(model_type), model_type, image, text, state)
         bpe_text = format_bpe_display(bpe)
-        return image, vis, bpe_text, slider_max_val
+        index_slider.update(visible=True, maximum=slider_max_val, value=0)
+        prev_btn.update(visible=True)
+        next_btn.update(visible=True)
+        return image, vis, bpe_text
+
+    def update_index(change, state):
+        state['current_index'] = max(0, min(len(state['current_vis']) - 1, state['current_index'] + change))
+        return state['current_vis'][state['current_index']], format_bpe_display(state['current_bpe'][state['current_index']])
+
+    def format_bpe_display(bpe):
+        return f"<div style='text-align:center; font-size:20px;'><strong>Current BPE: <span style='color:red;'>{bpe}</span></strong></div>"
 
-    
     run_btn.click(
         on_run_clicked,
-        inputs=[model_type, image_input, text_input],
-        outputs=[orig_img, heatmap, bpe_display, index_slider],
-    ).then(
-        lambda max_val: (gr.update(visible=True), gr.update(visible=True, maximum=max_val, value=0), gr.update(visible=True), gr.update(visible=True)),
-        inputs=index_slider,
-        outputs=[prev_btn, index_slider, next_btn, bpe_display],
+        inputs=[model_type, image_input, text_input, state],
+        outputs=[orig_img, heatmap, bpe_display],
     )
     
     prev_btn.click(
-        lambda: (*update_index(-1), current_index),
+        lambda state: (*update_index(-1, state), state['current_index']),
+        inputs=[state],
         outputs=[heatmap, bpe_display, index_slider]
     )
     
     next_btn.click(
-        lambda: (*update_index(1), current_index),
+        lambda state: (*update_index(1, state), state['current_index']),
+        inputs=[state],
         outputs=[heatmap, bpe_display, index_slider]
     )
 
 
     index_slider.change(
-            update_slider_index,
-            inputs=index_slider,
+            lambda x, state: update_slider_index(x, state),
+            inputs=[index_slider, state],
             outputs=[heatmap, bpe_display]
         )
 
-
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()