Update app.py

app.py

@@ -7,10 +7,10 @@ import pickle
 from insightface.app import FaceAnalysis
 from concurrent.futures import ThreadPoolExecutor
 
-# Initialize your models only once
-app = FaceAnalysis(name='buffalo_l')
-app.prepare(ctx_id=0, det_size=(640, 640))
-swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), 'inswapper_128.onnx'), download=False)
+# Initialize your models to None
+app = None
+swapper = None
+data_images = None
 
 # Load pickle file once and keep it in memory
 @st.cache_data
@@ -18,7 +18,21 @@ def load_data_images():
  with open(os.path.join(os.getcwd(), 'data_images.pkl'), 'rb') as file:
  return pickle.load(file)
 
-data_images = load_data_images()
+# Lazy initialization of data_images
+def get_data_images():
+ global data_images
+ if data_images is None:
+ data_images = load_data_images()
+ return data_images
+
+# Lazy initialization for models
+def prepare_models():
+ global app, swapper
+ if app is None:
+ app = FaceAnalysis(name='buffalo_l')
+ app.prepare(ctx_id=0, det_size=(640, 640))
+ if swapper is None:
+ swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), 'inswapper_128.onnx'), download=False)
 
 def face_swapper(image_background, image_customer):
  face_customer = app.get(image_customer)[0]
@@ -30,7 +44,8 @@ def face_swapper(image_background, image_customer):
  return image_background
 
 def process(image):
- images_background_encoding, images_background_contents = data_images['encodings'], data_images['content']
+ prepare_models() # Prepare models only when required
+ images_background_encoding, images_background_contents = get_data_images()['encodings'], get_data_images()['content']
  image_loaded = face_recognition.load_image_file(image)
  face_encoding = face_recognition.face_encodings(image_loaded)[0]
  face_distances = face_recognition.face_distance(images_background_encoding, face_encoding)
@@ -58,14 +73,15 @@ else:
 
 if uploaded_file is not None:
  bytes_data = uploaded_file.getvalue()
- if option=='File':
+ if option == 'File':
  st.image(uploaded_file)
- 
+
  if st.button('Process'):
- with ThreadPoolExecutor() as executor:
- future = executor.submit(process, uploaded_file)
- image_output = future.result()
- st.image(image_output)
+ with st.spinner('Processing...'):
+ with ThreadPoolExecutor() as executor:
+ future = executor.submit(process, uploaded_file)
+ image_output = future.result()
+ st.image(image_output)
 
 if image_output is not None:
  image_output_to_download = cv2.cvtColor(image_output, cv2.COLOR_BGR2RGB)