Update app.py

app.py

@@ -12,43 +12,44 @@ peft_model_ckp = "./model/blip2_peft" #PEFT model checkpoint path
 sample_img_path = "./sample_images/"
 
 #init_model_required = True
-#processor = None
-#model = None
 
-#def init_model():
+def init_model():
 
     #if init_model_required:
 
-#Preprocess input 
-processor = Blip2Processor.from_pretrained(preprocess_ckp)
+    #Preprocess input 
+    processor = Blip2Processor.from_pretrained(preprocess_ckp)
 
-#Model   
-#Inferance on GPU device. Will give error in CPU system, as "load_in_8bit" is an setting of bitsandbytes library and only works for GPU
-#model = Blip2ForConditionalGeneration.from_pretrained(base_model_ckp, load_in_8bit = True, device_map = "auto") 
+    #Model   
+    #Inferance on GPU device. Will give error in CPU system, as "load_in_8bit" is an setting of bitsandbytes library and only works for GPU
+    #model = Blip2ForConditionalGeneration.from_pretrained(base_model_ckp, load_in_8bit = True, device_map = "auto") 
 
-#Inferance on CPU device
-model = Blip2ForConditionalGeneration.from_pretrained(base_model_ckp) 
+    #Inferance on CPU device
+    model = Blip2ForConditionalGeneration.from_pretrained(base_model_ckp) 
 
-model = PeftModel.from_pretrained(model, peft_model_ckp)
+    model = PeftModel.from_pretrained(model, peft_model_ckp)
 
         #init_model_required = False
 
+    return processor, model
     
 
 def main():
 
     st.title("Fashion Image Caption using BLIP2")
 
-    #init_model()
+    #processor, model = init_model()
 
     #Select few sample images for the catagory of cloths
-    #option = st.selectbox('Select from sample an images', ('cap', 'tee', 'dress'), index = -1)
-    #st.text("OR")
+    option = st.selectbox('Select from sample an images', ('cap', 'tee', 'dress'), index = -1)
+    st.text("OR")
     file_name = st.file_uploader("Upload an image")
+    st.text(option)
+    
+    """
+    if file_name is None and option is not None:
 
-    #if file_name is None and option is not None:
-
-    #    file_name = os.path.join(sample_img_path, option)
+        file_name = os.path.join(sample_img_path, option)
 
     if file_name is not None:
 
@@ -74,6 +75,7 @@ def main():
         #Output the predict text
         caption_text.header("Generated Caption")
         caption_text.text(generated_caption)
+    """
 
 if __name__ == "__main__":
     main()