Update app.py

app.py

@@ -3,8 +3,9 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoProcessor
 import torch
 from PIL import Image
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Prompt templates
 user_prompt = '<|user|>\n'
 assistant_prompt = '<|assistant|>\n'
 prompt_suffix = "<|end|>\n"
@@ -30,13 +31,12 @@ def run_example(image, text_input=None, model_id=model_name):
     prompt = f"{user_prompt}<|image_1|>\n{text_input}{prompt_suffix}{assistant_prompt}"
     image = Image.fromarray(image).convert("RGB")
 
-    inputs = processor(prompt, image, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        generate_ids = model.generate(
-            **inputs,
-            max_new_tokens=1000,
-            eos_token_id=processor.tokenizer.eos_token_id
-        )
+    inputs = processor(prompt, image, return_tensors="pt").to("cuda:0")
+    generate_ids = model.generate(
+        **inputs,
+        max_new_tokens=1000,
+        eos_token_id=processor.tokenizer.eos_token_id
+    )
     generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
     response = processor.batch_decode(
         generate_ids,
@@ -53,7 +53,7 @@ css = """
   }
 """
 
-with gr.Blocks(css=css, title="Phi-3.5 Vision Instruct Demo") as demo:
+with gr.Blocks(css=css) as demo:
     gr.Markdown("## Phi-3.5 Vision Instruct Demo with Example Inputs")
 
     with gr.Tab(label="Phi-3.5 Input"):
@@ -82,12 +82,7 @@ with gr.Blocks(css=css, title="Phi-3.5 Vision Instruct Demo") as demo:
             examples_per_page=3
         )
 
-        submit_btn.click(
-            fn=run_example,
-            inputs=[input_img, text_input, model_selector],
-            outputs=output_text
-        )
+        submit_btn.click(run_example, [input_img, text_input, model_selector], [output_text])
 
-# Queue and launch without share (handled automatically on Spaces)
 demo.queue()
-demo.launch()
+demo.launch()