Update app.py

app.py

@@ -1,9 +1,13 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
 import gradio as gr
 from threading import Thread
 from PIL import Image
+import subprocess
+
+# Install flash-attention
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # Constants
 TITLE = "<h1><center>Phi 3.5 Multimodal (Text + Vision)</center></h1>"
@@ -14,23 +18,29 @@ 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}")
+
+# Quantization config for text model
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
 
 # Load models and tokenizers
 text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
 text_model = AutoModelForCausalLM.from_pretrained(
     TEXT_MODEL_ID,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
-    low_cpu_mem_usage=True
+    quantization_config=quantization_config
 )
 
 vision_model = AutoModelForCausalLM.from_pretrained(
     VISION_MODEL_ID, 
     trust_remote_code=True, 
-    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
+    torch_dtype="auto", 
+    attn_implementation="flash_attention_2"
 ).to(device).eval()
 
 vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
@@ -45,7 +55,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(device)
+    input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(text_model.device)
     streamer = TextIteratorStreamer(text_tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(