Update demo.py

demo.py

@@ -6,12 +6,12 @@ import torch
 model_id = "MargiPandya/blip2-bdd100k"
 processor = AutoProcessor.from_pretrained(model_id)
 config = PeftConfig.from_pretrained(model_id)
-model = Blip2ForConditionalGeneration.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map="auto")
+model = Blip2ForConditionalGeneration.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map="cpu")
 model = PeftModel.from_pretrained(model, model_id)
 
 
 def caption_image(image):
-    inputs = processor(images=image, return_tensors="pt").to("cuda", torch.float16)
+    inputs = processor(images=image, return_tensors="pt").to("cpu", torch.float16)
     generated_ids = model.generate(
           **inputs,
           max_new_tokens=64,