Update app.py

app.py

@@ -1,25 +1,34 @@
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 import torch
 from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 import os
+os.makedirs("./offload", exist_ok=True)
+from accelerate import infer_auto_device_map
 
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
 
-# BLIP-2 for image captioning
-blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
+
+# For BLIP-2
 blip_model = Blip2ForConditionalGeneration.from_pretrained(
     "Salesforce/blip2-opt-2.7b",
     torch_dtype=torch.float16,
-    device_map="auto"
+    device_map="auto",
+    offload_folder="./offload",
+    no_split_module_classes=["Blip2QFormerModel"]
 )
 
-# Phi-3 for text generation
+# For Phi-3
 phi_model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct",
     trust_remote_code=True,
     device_map="auto",
     torch_dtype=torch.float16,
-    token=HF_TOKEN
+    offload_folder="./offload",
+    no_split_module_classes=["PhiDecoderLayer"],
+    load_in_4bit=True  # Add 4-bit quantization
 )
 phi_tokenizer = AutoTokenizer.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct",