Update app.py

app.py

@@ -176,6 +176,7 @@ def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")  # Use PNG for compatibility with RGBA
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
+    
 @spaces.GPU
 def rm_background(image):
     im = load_img(image, output_type="pil")
@@ -204,7 +205,8 @@ def remove_background(image):
     else:
         raise TypeError("Unsupported image type")
     return output
-
+    
+@spaces.GPU
 def detect_and_segment_persons(image, clothes):
     img = np.array(image)
     img = img[:, :, ::-1]  # RGB -> BGR
@@ -272,5 +274,14 @@ def get_image(image_id):
     except FileNotFoundError:
         return jsonify({'error': 'Image not found'}), 404
 
+def initialize_model_in_subprocess():
+    # Charger le modèle avec le GPU ici
+    model = torch.hub.load("pytorch/vision:v0.10.0", "resnet18", pretrained=True)
+    model = model.to("cuda")
+    return model
+        
+
 if __name__ == "__main__":
+    with multiprocessing.Pool(1) as pool:
+        model = pool.apply(initialize_model_in_subprocess)
     app.run(debug=True, host="0.0.0.0", port=7860)