Update app.py

app.py

@@ -1,13 +1,9 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer
 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>"
@@ -19,28 +15,23 @@ VISION_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# 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.bfloat16,
-    device_map="auto",
-    quantization_config=quantization_config
+    torch_dtype=torch.float32 if device == "cpu" else torch.float16,
+    device_map="auto" if device == "cuda" else None,
+    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="auto", 
-    attn_implementation="flash_attention_2"
+    torch_dtype=torch.float32 if device == "cpu" else torch.float16,
+    low_cpu_mem_usage=True
 ).to(device).eval()
 
 vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
@@ -125,4 +116,5 @@ 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()