:zap: Fix version

README.md

@@ -4,7 +4,7 @@ emoji: 🔍
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 4.16.0
+sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -1,8 +1,7 @@
 import gradio as gr
-import spaces
 import torch
 import numpy as np
-
+import spaces
 from ola_vlm.constants import DEFAULT_IMAGE_TOKEN
 
 from ola_vlm.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
@@ -23,6 +22,14 @@ import math
 from transformers import TextIteratorStreamer
 from threading import Thread
 
+import subprocess
+# Install flash attention, skipping CUDA build if necessary
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
+
 def make_grid(pil_images, layer_indices=None):
     new_images = []
     new_captions = []
@@ -242,48 +249,51 @@ def regenerate(state, image_process_mode):
     state.skip_next = False
     return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
-@spaces.GPU
-def get_interm_outs(state):
-    prompt = state.get_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")
-            
-            #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)
+# @spaces.GPU
+# def get_interm_outs(state):
+    
 
-            if type(inp_images) is list:
-                inp_images = [image.to(model.device, dtype=torch.float16) for image in images]
+@spaces.GPU
+def generate(state, temperature, top_p, max_output_tokens, is_inter=False):
+    if is_inter:
+        prompt = state.get_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")
+                
+                #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 = inp_images.to(model.device, dtype=torch.float16)
+                inp_images = None
+                image_sizes = None
+            image_args = {"images": inp_images, "image_sizes": image_sizes}
         else:
             inp_images = None
-            image_sizes = None
-        image_args = {"images": inp_images, "image_sizes": image_sizes}
-    else:
-        inp_images = None
-        image_args = {}
+            image_args = {}
 
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+        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)
+        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)
 
-    return depth_outs, seg_outs, gen_outs
+        return depth_outs, seg_outs, gen_outs
 
-@spaces.GPU
-def generate(state, temperature, top_p, max_output_tokens):
     prompt = state.get_prompt()
     images = state.get_images(return_pil=True)
     #prompt, image_args = process_image(prompt, images)
@@ -439,9 +449,9 @@ 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(
-        get_interm_outs,
+        generate,
         [state],
-        [depth_box, seg_box, gen_box],
+        [depth_box, seg_box, gen_box, True],
     )
 
     clear_btn.click(