Model in memory

app.py

@@ -66,27 +66,33 @@ st.image(image, caption='Your target document')
 with st.spinner(f'Processing the document ...'):
         pre_trained_model = "unstructuredio/chipper-fast-fine-tuning"
         processor = DonutProcessor.from_pretrained(pre_trained_model)
-        model = VisionEncoderDecoderModel.from_pretrained(pre_trained_model)
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        from huggingface_hub import hf_hub_download
+        if 'model' in st.session_state:
+            model = st.session_state['model']
+        else:
+            model = VisionEncoderDecoderModel.from_pretrained(pre_trained_model)
 
-        lm_head_file = hf_hub_download(
-            repo_id=pre_trained_model, filename="lm_head.pth"
-        )
+            from huggingface_hub import hf_hub_download
 
-        rank = 128
-        model.decoder.lm_head = nn.Sequential(
-            nn.Linear(model.decoder.lm_head.weight.shape[1], rank, bias=False),
-            nn.Linear(rank, rank, bias=False),
-            nn.Linear(rank, model.decoder.lm_head.weight.shape[0], bias=True),
-        )
+            lm_head_file = hf_hub_download(
+                repo_id=pre_trained_model, filename="lm_head.pth"
+            )
 
-        model.decoder.lm_head.load_state_dict(torch.load(lm_head_file))
+            rank = 128
+            model.decoder.lm_head = nn.Sequential(
+                nn.Linear(model.decoder.lm_head.weight.shape[1], rank, bias=False),
+                nn.Linear(rank, rank, bias=False),
+                nn.Linear(rank, model.decoder.lm_head.weight.shape[0], bias=True),
+            )
+
+            model.decoder.lm_head.load_state_dict(torch.load(lm_head_file))
 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        model.eval()
-        model.to(device)
+            model.eval()
+            model.to(device)
+            st.session_state['model'] = model
 
 st.info(f'Parsing document')
 parsed_info = run_prediction(image.convert("RGB"), model, processor, prompt)