Update app.py

app.py

@@ -4,21 +4,43 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel, PeftConfig
 from PIL import Image
 import torchvision.datasets as datasets
+import numpy as np
+import os
 
 def load_model():
-    # Load base Phi model
+    # Create offload directory
+    os.makedirs("offload", exist_ok=True)
+    
+    # Configure device map for memory efficiency
+    device_map = {
+        'base_model.model.model.embed_tokens': 0,
+        'base_model.model.model.layers.0': 0,
+        'base_model.model.model.layers.1': 0,
+        'base_model.model.model.layers.2': 0,
+        'base_model.model.model.layers.3': 0,
+        'base_model.model.model.layers.4': 'cpu',
+        'base_model.model.model.layers.5': 'cpu',
+        'base_model.model.model.layers.6': 'cpu',
+        'base_model.model.model.layers.7': 'cpu',
+        'base_model.model.model.layers.8': 'cpu',
+        'base_model.model.model.norm': 'cpu',
+        'base_model.model.lm_head': 0,
+    }
+    
     base_model = AutoModelForCausalLM.from_pretrained(
         "microsoft/Phi-3-mini-4k-instruct",
         trust_remote_code=True,
-        device_map="auto",
-        torch_dtype=torch.float32
+        device_map=device_map,  # Use custom device map
+        torch_dtype=torch.float32,
+        attn_implementation='eager',
+        offload_folder="offload"
     )
     
-    # Load our fine-tuned LoRA adapter
     model = PeftModel.from_pretrained(
         base_model,
-        "jatingocodeo/phi-vlm",  # Your uploaded model
-        device_map="auto"
+        "jatingocodeo/phi-vlm",
+        device_map=device_map,
+        offload_folder="offload"
     )
     
     tokenizer = AutoTokenizer.from_pretrained("jatingocodeo/phi-vlm")
@@ -30,13 +52,16 @@ def generate_description(image, model, tokenizer):
     if image.mode != "RGB":
         image = image.convert("RGB")
     
-    # Resize image to match training size
+    # Resize image to match training size (32x32)
     image = image.resize((32, 32))
     
-    # Prepare prompt
-    prompt = """Below is an image. Please describe it in detail.
+    # Convert image to tensor and normalize
+    image_tensor = torch.FloatTensor(np.array(image)).permute(2, 0, 1) / 255.0
+    
+    # Prepare prompt with image tensor
+    prompt = f"""Below is an image. Please describe it in detail.
 
-Image: <image>
+Image: {image_tensor}
 Description: """
     
     # Tokenize input
@@ -51,7 +76,8 @@ Description: """
     # Generate description
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,
+            input_ids=inputs.input_ids,
+            attention_mask=inputs.attention_mask,
             max_new_tokens=100,
             temperature=0.7,
             do_sample=True,