Update app.py

app.py

@@ -9,7 +9,7 @@ from datetime import datetime
 import numpy as np
 import os
 
-HF_TOKEN = os.environ['HF_TOKEN']
+
 # subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # models = {
@@ -21,8 +21,10 @@ def array_to_image_path(image_array):
     img = Image.fromarray(np.uint8(image_array))
     
     # Generate a unique filename using timestamp
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"image_{timestamp}.png"
+    # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    # filename = f"image_{timestamp}.png" # comment this, only save 1 image int the local path
+
+    filename = "image_to_inference.png"
     
     # Save the image
     img.save(filename)
@@ -33,21 +35,21 @@ def array_to_image_path(image_array):
     return full_path
     
 models = {
-    "Qwen/Qwen2-VL-7B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained(
-        "Qwen/Qwen2-VL-7B-Instruct", 
+    "Qwen/Qwen2-VL-2B-Instruct": Qwen2VLForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2-VL-2B-Instruct", 
         trust_remote_code=True, 
-        token=HF_TOKEN,
-        torch_dtype=torch.bfloat16,
-        # attn_implementation="flash_attention_2"
+        # load_in_4bit=True,
+        # attn_implementation="flash_attention_2",
+        torch_dtype=torch.bfloat16
     ).cuda().eval()
 
 }
 
 processors = {
-    "Qwen/Qwen2-VL-7B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True)
+    "Qwen/Qwen2-VL-2B-Instruct": AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True)
 }
 
-DESCRIPTION = "[Qwen2-VL-7B Demo](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct)"
+DESCRIPTION = "[Qwen2-VL-2B Demo](https://huggingface.co/Qwen/Qwen2-VL-2B-Instruct)"
 
 kwargs = {}
 kwargs['torch_dtype'] = torch.bfloat16
@@ -57,7 +59,7 @@ assistant_prompt = '<|assistant|>\n'
 prompt_suffix = "<|end|>\n"
 
 @spaces.GPU
-def run_example(image, text_input=None, model_id="Qwen/Qwen2-VL-7B-Instruct"):
+def run_example(image, text_input=None, model_id="Qwen/Qwen2-VL-2B-Instruct"):
     image_path = array_to_image_path(image)
     
     print(image_path)
@@ -114,11 +116,11 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
-    with gr.Tab(label="Qwen2-VL-7B Input"):
+    with gr.Tab(label="Qwen2-VL-2B Input"):
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="Input Picture")
-                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="Qwen/Qwen2-VL-7B-Instruct")
+                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value="Qwen/Qwen2-VL-2B-Instruct")
                 text_input = gr.Textbox(label="Question")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():