Update app.py

app.py

@@ -14,23 +14,22 @@ TEXT_MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
 VISION_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
 # Load models and tokenizers
 text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
 text_model = AutoModelForCausalLM.from_pretrained(
     TEXT_MODEL_ID,
-    torch_dtype=torch.float32 if device == "cpu" else torch.float16,
-    device_map="auto" if device == "cuda" else None,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto",
     low_cpu_mem_usage=True
 )
 
-if device == "cuda":
-    text_model = text_model.half()  # Convert to half precision if on GPU
-
 vision_model = AutoModelForCausalLM.from_pretrained(
     VISION_MODEL_ID, 
     trust_remote_code=True, 
-    torch_dtype=torch.float32 if device == "cpu" else torch.float16,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    attn_implementation="flash_attention_2" if device == "cuda" else None,
     low_cpu_mem_usage=True
 ).to(device).eval()
 
@@ -46,7 +45,7 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
         ])
     conversation.append({"role": "user", "content": message})
 
-    input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(text_model.device)
+    input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(text_tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
@@ -116,5 +115,4 @@ with gr.Blocks() as demo:
         vision_submit_btn.click(process_vision_query, [vision_input_img, vision_text_input], [vision_output_text])
 
 if __name__ == "__main__":
-    print(f"Running on device: {device}")
     demo.launch()