Create gradio_ootd.py

run/gradio_ootd.py

@@ -0,0 +1,118 @@
+from flask import Flask, request, jsonify ,send_file
+from PIL import Image
+import base64
+import io 
+import random
+import uuid
+import numpy as np
+import spaces
+import torch
+
+from utils_ootd import get_mask_location
+
+PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute()
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from preprocess.openpose.run_openpose import OpenPose
+from preprocess.humanparsing.run_parsing import Parsing
+from ootd.inference_ootd_hd import OOTDiffusionHD
+from ootd.inference_ootd_dc import OOTDiffusionDC
+
+
+openpose_model_hd = OpenPose(0)
+parsing_model_hd = Parsing(0)
+ootd_model_hd = OOTDiffusionHD(0)
+
+openpose_model_dc = OpenPose(1)
+parsing_model_dc = Parsing(1)
+ootd_model_dc = OOTDiffusionDC(1)
+
+
+category_dict = ['upperbody', 'lowerbody', 'dress']
+category_dict_utils = ['upper_body', 'lower_body', 'dresses']
+
+
+
+# Créer une instance FastAPI
+app = Flask(__name__)
+
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+@spaces.GPU
+def process_hd(vton_img, garm_img,category, n_samples, n_steps, image_scale, seed):
+    model_type = 'hd'
+    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')
+        
+        garm_img = Image.open(garm_img).resize((768, 1024))
+        vton_img = Image.open(vton_img).resize((768, 1024))
+        keypoints = openpose_model_hd(vton_img.resize((384, 512)))
+        model_parse, _ = parsing_model_hd(vton_img.resize((384, 512)))
+
+        mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints)
+        mask = mask.resize((768, 1024), Image.NEAREST)
+        mask_gray = mask_gray.resize((768, 1024), Image.NEAREST)
+        
+        masked_vton_img = Image.composite(mask_gray, vton_img, mask)
+
+        images = ootd_model_hd(
+            model_type=model_type,
+            category=category_dict[category],
+            image_garm=garm_img,
+            image_vton=masked_vton_img,
+            mask=mask,
+            image_ori=vton_img,
+            num_samples=n_samples,
+            num_steps=n_steps,
+            image_scale=image_scale,
+            seed=seed,
+        )
+
+    return images
+
+
+@app.get("/")
+def root():
+    return "Welcome to the Fashion OOTDiffusion API "
+
+# Route pour récupérer l'image générée
+@app.route('/api/get_image/<image_id>', methods=['GET'])
+def get_image(image_id):
+    # Construire le chemin complet de l'image
+    image_path = image_id  # Assurez-vous que le nom de fichier correspond à celui que vous avez utilisé lors de la sauvegarde
+
+    # Renvoyer l'image
+    try:
+        return send_file(image_path, mimetype='image/png')
+    except FileNotFoundError:
+        return jsonify({'error': 'Image not found'}), 404
+
+# Route pour l'API REST
+@app.route('/api/run', methods=['POST'])
+def run():
+    data = request.json
+    print(data)
+    vton_img = data['vton_img']
+    garm_img = data['garm_img']
+    category = data['category']
+    n_samples = data['n_samples']
+    n_steps = data['n_steps']
+    image_scale = data['image_scale']
+    seed = data['seed']
+    result = process_hd(vton_img, garm_img,category, n_samples, n_steps, image_scale, seed)
+    return jsonify({'out': result})
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)