Upload BLIPIntepret.py

BLIPIntepret.py

@@ -7,14 +7,14 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 print(device)
 
 def init_BLIP(device):
+    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
     if device == 'cuda':
-    	bit_load = True
+        model = Blip2ForConditionalGeneration.from_pretrained(
+         "Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16, device_map = 'auto')
     else:
-    	bit_load = False
-    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-    model = Blip2ForConditionalGeneration.from_pretrained(
-        "Salesforce/blip2-opt-2.7b", load_in_8bit= bit_load,torch_dtype=torch.float16, device_map = 'auto'
-    )
+        print('Using CPU model')
+        model = Blip2ForConditionalGeneration.from_pretrained( "Salesforce/blip2-opt-2.7b",device_map={"": device}, torch_dtype=torch.float32,low_cpu_mem_usage=True)
+
     model.eval()
     if torch.__version__ >= "2":
         model = torch.compile(model)
@@ -33,8 +33,10 @@ def infer_BLIP2(model,processor,image,device):
         "Question: What emotion does the person or animal in the image feel? Answer:",
         ]
     for prompt in prompts:
-        inputs = processor(images=image, text=prompt, return_tensors="pt").to(device, torch.float16)
-
+        if device == 'cuda':
+            inputs = processor(images=image, text=prompt, return_tensors="pt").to(device, torch.float16)
+        else:
+            inputs = processor(images=image, text=prompt, return_tensors="pt")
         generated_ids = model.generate(**inputs)
         generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
         outputs+= prompt+generated_text+' '