Update app.py

app.py

@@ -23,27 +23,27 @@ map_sampleid_name = {
                     'tee': 'e2d8637a-5478-429d-a2a8-3d5859dbc64d.jpeg',
                     'bracelet': 'e78518ac-0f54-4483-a233-fad6511f0b86.jpeg'
                     }
-#init_model_required = True
+init_model_required = True
 
-def init_model():
+def init_model(init_model_required):
 
-    #if init_model_required:
+    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
+        init_model_required = False
 
-    return processor, model
+    return processor, model, init_model_required
 
 #def main():
 
@@ -63,6 +63,7 @@ st.text("Or")
 file_name = st.file_uploader(label = "Upload an image", accept_multiple_files = False)
 
 btn_click = st.button('Generate')
+st.caption("Application deployed on CPU basic with 16GB RAM")
 
 if btn_click:
 
@@ -79,7 +80,7 @@ if btn_click:
     if image is not None:
 
         with st.spinner('Initializing model...'):
-            processor, model = init_model()
+            processor, model, init_model_required = init_model(init_model_required)
 
         with st.spinner('Generating Caption...'):
 
@@ -104,7 +105,6 @@ if btn_click:
             caption_text.header("Generated Caption")
             caption_text.text(generated_caption)
 
-st.caption("Application deployed on CPU basic with 16GB RAM")
 
 #if __name__ == "__main__":
 #   main()