Update app.py

app.py

@@ -5,6 +5,7 @@ from io import BytesIO
 import torch
 import base64
 import io 
+import logging
 import gradio as gr
 import numpy as np
 import spaces
@@ -149,7 +150,7 @@ def start_tryon_full_body(tops_img, bottoms_img, model_parse_tops, model_parse_b
     human_img_arg_tops = _apply_exif_orientation(human_img_orig_tops.resize((384, 512)))
     human_img_arg_tops = convert_PIL_to_numpy(human_img_arg_tops, format="BGR")
 
-    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
+    args = apply_net.create_argument_parser().parse_args(['show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'])
     pose_img_tops = args.func(args, human_img_arg_tops)
     pose_img_tops = pose_img_tops[:, :, ::-1]
     pose_img_tops = Image.fromarray(pose_img_tops).resize((768, 1024))
@@ -367,16 +368,24 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
         return images[0], mask_gray
 
 
-# Fonction pour décoder une image encodée en base64 en objet PIL.Image.Image
-def decode_image_from_base64(image_data):
-    image_data = base64.b64decode(image_data)
-    image = Image.open(io.BytesIO(image_data))
-    return image
-
-def encode_image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+def decode_image_from_base64(base64_str):
+    try:
+        img_data = base64.b64decode(base64_str)
+        img = Image.open(BytesIO(img_data))
+        return img
+    except Exception as e:
+        logging.error(f"Error decoding image: {e}")
+        raise
+
+def encode_image_to_base64(img):
+    try:
+        buffered = BytesIO()
+        img.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return img_str
+    except Exception as e:
+        logging.error(f"Error encoding image: {e}")
+        raise
 
 @app.route('/tryon', methods=['POST'])
 def tryon():
@@ -409,41 +418,45 @@ def tryon():
 
 @app.route('/tryon-full', methods=['POST'])
 def tryon_full():
-    data = request.json
-
-    # Decode images
-    tops_image = decode_image_from_base64(data['tops_image'])
-    bottoms_image = decode_image_from_base64(data['bottoms_image'])
-    model_parse_tops = decode_image_from_base64(data['model_parse_tops'])
-    model_parse_bottoms = decode_image_from_base64(data['model_parse_bottoms'])
-
-    # Retrieve additional parameters
-    is_checked = data.get('use_auto_mask', True)
-    is_checked_crop = data.get('use_auto_crop', False)
-    denoise_steps = int(data.get('denoise_steps', 30))
-    seed = int(data.get('seed', 42))
-
-    # Call the start_tryon_full_body function
-    output_image, mask_image = start_tryon_full_body(
-        tops_image,
-        bottoms_image,
-        model_parse_tops,
-        model_parse_bottoms,
-        is_checked,
-        is_checked_crop,
-        denoise_steps,
-        seed
-    )
-
-    # Convert output image to base64
-    output_base64 = encode_image_to_base64(output_image)
-    mask_base64 = encode_image_to_base64(mask_image)
-
-    return jsonify({
-        'output_image': output_base64,
-        'mask_image': mask_base64
-    })
-
+    try:
+        data = request.json
+
+        # Decode images
+        tops_image = decode_image_from_base64(data['tops_image'])
+        bottoms_image = decode_image_from_base64(data['bottoms_image'])
+        model_parse_tops = decode_image_from_base64(data['model_parse_tops'])
+        model_parse_bottoms = decode_image_from_base64(data['model_parse_bottoms'])
+
+        # Retrieve additional parameters
+        is_checked = data.get('use_auto_mask', True)
+        is_checked_crop = data.get('use_auto_crop', False)
+        denoise_steps = int(data.get('denoise_steps', 30))
+        seed = int(data.get('seed', 42))
+
+        # Call the start_tryon_full_body function
+        output_image, mask_image = start_tryon_full_body(
+            tops_image,
+            bottoms_image,
+            model_parse_tops,
+            model_parse_bottoms,
+            is_checked,
+            is_checked_crop,
+            denoise_steps,
+            seed
+        )
+
+        # Convert output image to base64
+        output_base64 = encode_image_to_base64(output_image)
+        mask_base64 = encode_image_to_base64(mask_image)
+
+        return jsonify({
+            'output_image': output_base64,
+            'mask_image': mask_base64
+        })
+    
+    except Exception as e:
+        logging.error(f"Error in /tryon-full: {e}")
+        return jsonify({'error': str(e)}), 200
 
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)