Update app.py

app.py

@@ -1,35 +1,36 @@
 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
+
+# Initialize environment
 os.makedirs("./offload", exist_ok=True)
-from accelerate import infer_auto_device_map
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
+# Memory optimization
+torch.cuda.empty_cache()
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
-
-# For BLIP-2
+# Load BLIP-2
+blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
 blip_model = Blip2ForConditionalGeneration.from_pretrained(
     "Salesforce/blip2-opt-2.7b",
     torch_dtype=torch.float16,
-    device_map="auto",
-    offload_folder="./offload",
-    no_split_module_classes=["Blip2QFormerModel"]
-)
+    device_map="auto"
+).eval()
 
-# For Phi-3
+# Load Phi-3
 phi_model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct",
     trust_remote_code=True,
     device_map="auto",
     torch_dtype=torch.float16,
-    offload_folder="./offload",
-    no_split_module_classes=["PhiDecoderLayer"],
-    load_in_4bit=True  # Add 4-bit quantization
-)
+    load_in_4bit=True,
+    token=HF_TOKEN
+).eval()
 phi_tokenizer = AutoTokenizer.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct",
     token=HF_TOKEN