Update app.py

app.py

@@ -10,17 +10,9 @@ import io
 import logging
 import gradio as gr
 import numpy as np
+import spaces
 import uuid
 import random
-if os.environ.get("SPACES_ZERO_GPU") is not None:
-    import spaces
-else:
-    class spaces:
-        @staticmethod
-        def GPU(func):
-            def wrapper(*args, **kwargs):
-                return func(*args, **kwargs)
-            return wrapper
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
@@ -74,7 +66,6 @@ UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="un
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
-
 # Préparation du pipeline Tryon
 pipe = TryonPipeline.from_pretrained(
     base_path,
@@ -98,7 +89,6 @@ tensor_transfrom = transforms.Compose([
     transforms.Normalize([0.5], [0.5]),
 ])
 
-@spaces.GPU
 def pil_to_binary_mask(pil_image, threshold=0):
     np_image = np.array(pil_image)
     grayscale_image = Image.fromarray(np_image).convert("L")
@@ -142,6 +132,7 @@ def save_image(img):
     img.save(unique_name, format="WEBP", lossless=True) 
     return unique_name
 
+@spaces.GPU
 def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed, categorie = 'upper_body'):
     device = "cuda"
     openpose_model.preprocessor.body_estimation.model.to(device)
@@ -238,7 +229,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
                     height=768,
                     width=768,
                     ip_adapter_image=garm_img.resize((768, 768)),
-                    guidance_scale=2,
+                    guidance_scale=5,
                 )[0]
 
     if is_checked_crop:
@@ -248,7 +239,7 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
     else:
         return images[0], mask_gray , mask
 
-@spaces.GPU        
+        
 @app.route('/tryon-v2', methods=['POST'])
 def tryon_v2():
 
@@ -295,7 +286,6 @@ def process_image(image_data):
     else:
         return decode_image_from_base64(image_data)  # Décode l'image base64
 
-@spaces.GPU
 @app.route('/tryon', methods=['POST'])
 def tryon():
     data = request.json
@@ -326,7 +316,6 @@ def tryon():
 
 
 # Route index
-@spaces.GPU
 @app.route('/', methods=['GET'])
 def index():
 
@@ -350,6 +339,7 @@ def get_image(image_id):
 		
 		
 
+@spaces.GPU 
 def generate_mask(human_img, categorie='upper_body'):
     device = "cuda"
     openpose_model.preprocessor.body_estimation.model.to(device)
@@ -373,7 +363,6 @@ def generate_mask(human_img, categorie='upper_body'):
         logging.error(f"Error generating mask: {e}")
         raise e
 
-@spaces.GPU
 @app.route('/generate_mask', methods=['POST'])
 def generate_mask_api():
     try:
@@ -398,7 +387,10 @@ def generate_mask_api():
         logging.error(f"Error generating mask: {e}")
         return jsonify({'error': str(e)}), 500
 
-
+@spaces.GPU
+def warmup_gpu():
+    logging.info("Warmup GPU called to register with Hugging Face Spaces.")
 
 if __name__ == "__main__":
+    warmup_gpu()
     app.run(debug=False, host="0.0.0.0", port=7860)