Update run/gradio_ootd.py

run/gradio_ootd.py

@@ -1,5 +1,4 @@
 from flask import Flask, request, jsonify
-import os
 import torch
 from PIL import Image, ImageOps
 
@@ -11,6 +10,7 @@ from ootd.inference_ootd_dc import OOTDiffusionDC
 
 app = Flask(__name__)
 
+# Charger les modèles une seule fois au démarrage de l'application
 openpose_model_hd = OpenPose(0)
 parsing_model_hd = Parsing(0)
 ootd_model_hd = OOTDiffusionHD(0)
@@ -19,6 +19,9 @@ openpose_model_dc = OpenPose(1)
 parsing_model_dc = Parsing(1)
 ootd_model_dc = OOTDiffusionDC(1)
 
+# Définir la configuration GPU
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 category_dict = ['upperbody', 'lowerbody', 'dress']
 category_dict_utils = ['upper_body', 'lower_body', 'dresses']
 
@@ -35,11 +38,12 @@ def process_hd():
     model_type = 'hd'
     category = 0 # 0:upperbody; 1:lowerbody; 2:dress
 
+    # Charger les modèles en mémoire GPU
     with torch.no_grad():
-        openpose_model_hd.preprocessor.body_estimation.model.to('cuda')
-        ootd_model_hd.pipe.to('cuda')
-        ootd_model_hd.image_encoder.to('cuda')
-        ootd_model_hd.text_encoder.to('cuda')
+        openpose_model_hd.preprocessor.body_estimation.model.to(device)
+        ootd_model_hd.pipe.to(device)
+        ootd_model_hd.image_encoder.to(device)
+        ootd_model_hd.text_encoder.to(device)
         
         garm_img = Image.open(garm_img).resize((768, 1024))
         vton_img = Image.open(vton_img).resize((768, 1024))
@@ -86,11 +90,12 @@ def process_dc():
     else:
         category = 2
 
+    # Charger les modèles en mémoire GPU
     with torch.no_grad():
-        openpose_model_dc.preprocessor.body_estimation.model.to('cuda')
-        ootd_model_dc.pipe.to('cuda')
-        ootd_model_dc.image_encoder.to('cuda')
-        ootd_model_dc.text_encoder.to('cuda')
+        openpose_model_dc.preprocessor.body_estimation.model.to(device)
+        ootd_model_dc.pipe.to(device)
+        ootd_model_dc.image_encoder.to(device)
+        ootd_model_dc.text_encoder.to(device)
         
         garm_img = Image.open(garm_img).resize((768, 1024))
         vton_img = Image.open(vton_img).resize((768, 1024))