Update app.py

app.py

@@ -251,44 +251,41 @@ def regenerate(state, image_process_mode):
 
 @spaces.GPU
 def get_interm_outs(state):
-    print("HERERERE")
-    print(state)
     prompt = state.get_prompt()
-    print(prompt)
     images = state.get_images(return_pil=True)
     #prompt, image_args = process_image(prompt, images)
 
-    # if images is not None and len(images) > 0:
-    #     if len(images) > 0:
-    #         if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
-    #             raise ValueError("Number of images does not match number of <image> tokens in prompt")
+    if images is not None and len(images) > 0:
+        if len(images) > 0:
+            if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
+                raise ValueError("Number of images does not match number of <image> tokens in prompt")
             
-    #         #images = [load_image_from_base64(image) for image in images]
-    #         image_sizes = [image.size for image in images]
-    #         inp_images = process_images(images, image_processor, model.config)
-
-    #         if type(inp_images) is list:
-    #             inp_images = [image.to(model.device, dtype=torch.float16) for image in images]
-    #         else:
-    #             inp_images = inp_images.to(model.device, dtype=torch.float16)
-    #     else:
-    #         inp_images = None
-    #         image_sizes = None
-    #     image_args = {"images": inp_images, "image_sizes": image_sizes}
-    # else:
-    #     inp_images = None
-    #     image_args = {}
-
-    # input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-
-    # interm_outs = model.get_visual_interpretations(
-    #             input_ids,
-    #             **image_args
-    #     )
+            #images = [load_image_from_base64(image) for image in images]
+            image_sizes = [image.size for image in images]
+            inp_images = process_images(images, image_processor, model.config)
+
+            if type(inp_images) is list:
+                inp_images = [image.to(model.device, dtype=torch.float16) for image in images]
+            else:
+                inp_images = inp_images.to(model.device, dtype=torch.float16)
+        else:
+            inp_images = None
+            image_sizes = None
+        image_args = {"images": inp_images, "image_sizes": image_sizes}
+    else:
+        inp_images = None
+        image_args = {}
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+
+    interm_outs = model.get_visual_interpretations(
+                input_ids,
+                **image_args
+        )
     
-    # depth_outs = get_depth_images(interm_outs, image_sizes[0]) 
-    # seg_outs =  get_seg_images(interm_outs, images[0])
-    # gen_outs = get_gen_images(interm_outs)
+    depth_outs = get_depth_images(interm_outs, image_sizes[0]) 
+    seg_outs =  get_seg_images(interm_outs, images[0])
+    gen_outs = get_gen_images(interm_outs)
 
     return images[0], images[0], images[0]
     
@@ -450,7 +447,7 @@ with gr.Blocks(title="OLA-VLM", theme=gr.themes.Default(), css=block_css) as dem
     btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
 
     inter_vis_btn.click(
-        generate,
+        get_interm_outs,
         [state],
         [depth_box, seg_box, gen_box],
     )