Update app.py

app.py

@@ -38,10 +38,13 @@ class UI:
                 # print(3)
                 # pixel_values = pixel_values.to(torch.float32)
 
-                print(4)
+                print(3)
                 generated_ids = model.to(device, torch.float32).generate(pixel_values=pixel_values, max_length=25)
+                
+                print(4)
                 generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
+                print(5, generated_caption)
                 st.write(generated_caption)
                 
             st.success("Here you go!")
@@ -55,10 +58,12 @@ class UI:
 class Model:
     def load_model(self):
         peft_model_id = "Shrey23/Image-Captioning"
-        config = PeftConfig.from_pretrained(peft_model_id)
+        # config = PeftConfig.from_pretrained(peft_model_id)
         global model
         global processor
-        model = Blip2ForConditionalGeneration.from_pretrained(config.base_model_name_or_path, low_cpu_mem_usage=True, torch_dtype=torch.float16) #, device_map="auto", load_in_8bit=True
+        PATH = "./model/"
+        local_files_only=True
+        model = Blip2ForConditionalGeneration.from_pretrained(PATH, local_files_only=True, low_cpu_mem_usage=True, torch_dtype=torch.float16) #, device_map="auto", load_in_8bit=True
         model = PeftModel.from_pretrained(model, peft_model_id)
         processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")