Update app.py

app.py

@@ -4,19 +4,29 @@ import bitsandbytes
 import accelerate
 import scipy
 from PIL import Image
-
+import torch.nn as nn
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, InstructBlipProcessor, InstructBlipForConditionalGeneration
 
 
 def load_caption_model(blip2=False, instructblip=True):
 
+
+model = YourModel()
+
+
     if blip2:
         processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16)
         model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16)
+        if torch.cuda.device_count() > 1:
+            model = nn.DataParallel(model)
+            model.to('cuda')
         #model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto")
         
     if instructblip:
         model = InstructBlipForConditionalGeneration.from_pretrained("Salesforce/instructblip-vicuna-7b", load_in_8bit=True,torch_dtype=torch.float16)
+        if torch.cuda.device_count() > 1:
+            model = nn.DataParallel(model)
+            model.to('cuda')
         processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b", load_in_8bit=True,torch_dtype=torch.float16)
 
     return model, processor