Update app.py

app.py

@@ -6,17 +6,17 @@ from torchvision.transforms.functional import InterpolationMode
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 
+# Constants for ImageNet preprocessing
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
 
 # Build the image transform
 def build_transform(input_size):
-    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
     transform = T.Compose([
         T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
         T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
         T.ToTensor(),
-        T.Normalize(mean=MEAN, std=STD)
+        T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
     ])
     return transform
 
@@ -60,10 +60,9 @@ path = 'OpenGVLab/InternVL2_5-78B'
 model = AutoModel.from_pretrained(
     path,
     torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=True,
-    use_flash_attn=True,
-    trust_remote_code=True
-).eval().cuda()
+    trust_remote_code=True,
+    device_map="auto"  # Use device map for efficient memory handling
+)
 
 tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)