Update app.py

app.py

@@ -1,54 +1,49 @@
-import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
 
-import requests
-import torch
-import gradio as gr
-from PIL import Image
-from transformers import AutoModelForCausalLM, AutoProcessor
+# Default: Load the model on the available device(s)
+model = Qwen2VLForConditionalGeneration.from_pretrained(
+    "OS-Copilot/OS-Atlas-Base-7B", torch_dtype="auto", device_map="auto"
+)
+processor = AutoProcessor.from_pretrained("OS-Copilot/OS-Atlas-Base-7B")
 
-model_id_or_path = "rhymes-ai/Aria"
+messages = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "image",
+                "image": "./exmaples/images/web_6f93090a-81f6-489e-bb35-1a2838b18c01.png",
+            },
+            {"type": "text", "text": "In this UI screenshot, what is the position of the element corresponding to the command \"switch language of current page\" (with bbox)?"},
+        ],
+    }
+]
 
-model = AutoModelForCausalLM.from_pretrained(model_id_or_path, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)
 
-processor = AutoProcessor.from_pretrained(model_id_or_path, trust_remote_code=True)
+# Preparation for inference
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+image_inputs, video_inputs = process_vision_info(messages)
+inputs = processor(
+    text=[text],
+    images=image_inputs,
+    videos=video_inputs,
+    padding=True,
+    return_tensors="pt",
+)
+inputs = inputs.to("cuda")
 
-@spaces.GPU
-def run():
-        
-    image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"
-    
-    image = Image.open(requests.get(image_path, stream=True).raw)
-    
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"text": None, "type": "image"},
-                {"text": "what is the image?", "type": "text"},
-            ],
-        }
-    ]
-    
-    text = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=text, images=image, return_tensors="pt")
-    inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
-        output = model.generate(
-            **inputs,
-            max_new_tokens=500,
-            stop_strings=["<|im_end|>"],
-            tokenizer=processor.tokenizer,
-            do_sample=True,
-            temperature=0.9,
-        )
-        output_ids = output[0][inputs["input_ids"].shape[1]:]
-        result = processor.decode(output_ids, skip_special_tokens=True)
+# Inference: Generation of the output
+generated_ids = model.generate(**inputs, max_new_tokens=128)
 
-with gr.Blocks() as demo:
-    btn = gr.Button("Run")
-    out = gr.Markdown()
-    btn.click(run, outputs=out)
+generated_ids_trimmed = [
+    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+]
 
-demo.queue().launch()
+output_text = processor.batch_decode(
+    generated_ids_trimmed, skip_special_tokens=False, clean_up_tokenization_spaces=False
+)
+print(output_text)
+# <|object_ref_start|>language switch<|object_ref_end|><|box_start|>(576,12),(592,42)<|box_end|><|im_end|>