Update run/gradio_ootd.py

run/gradio_ootd.py

@@ -1,18 +1,10 @@
-from flask import Flask, request, jsonify ,send_file
-import base64
-import io 
-import random
-import uuid
-import numpy as np
-import spaces
-import torch
+import gradio as gr
 import os
 from pathlib import Path
 import sys
 import torch
 from PIL import Image, ImageOps
 
-
 from utils_ootd import get_mask_location
 
 PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute()
@@ -36,43 +28,74 @@ ootd_model_dc = OOTDiffusionDC(1)
 category_dict = ['upperbody', 'lowerbody', 'dress']
 category_dict_utils = ['upper_body', 'lower_body', 'dresses']
 
-torch.cuda.empty_cache() 
 
+example_path = os.path.join(os.path.dirname(__file__), 'examples')
+model_hd = os.path.join(example_path, 'model/model_1.png')
+garment_hd = os.path.join(example_path, 'garment/03244_00.jpg')
+model_dc = os.path.join(example_path, 'model/model_8.png')
+garment_dc = os.path.join(example_path, 'garment/048554_1.jpg')
 
-# 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
+import spaces
+
+@spaces.GPU
+def process_hd(vton_img, garm_img, n_samples, n_steps, image_scale, seed):
+    model_type = 'hd'
+    category = 0 # 0:upperbody; 1:lowerbody; 2:dress
+
+    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)))
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
+        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,
+        )
 
-# 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
+    return images
 
 @spaces.GPU
-def process_hd(vton_img, garm_img,category, n_samples, n_steps, image_scale, seed):
+def process_dc(vton_img, garm_img, category, n_samples, n_steps, image_scale, seed):
     model_type = 'dc'
+    if category == 'Upper-body':
+        category = 0
+    elif category == 'Lower-body':
+        category = 1
+    else:
+        category =2
+
     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')
-
-        garm_img = decode_image_from_base64(garm_img).resize((768, 1024))
-        vton_img = decode_image_from_base64(vton_img).resize((768, 1024))
+        
+        garm_img = Image.open(garm_img).resize((768, 1024))
+        vton_img = Image.open(vton_img).resize((768, 1024))
         keypoints = openpose_model_dc(vton_img.resize((384, 512)))
         model_parse, _ = parsing_model_dc(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)
@@ -95,37 +118,8 @@ def process_hd(vton_img, garm_img,category, n_samples, n_steps, image_scale, see
     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
+block = gr.Interface(fn=process_hd, inputs=["image", "image", "number", "number", "number", "number"], outputs="image", title="OOTDiffusion Demo HD")
+block.launch()
 
-    # 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
-@spaces.GPU
-@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)
+block_dc = gr.Interface(fn=process_dc, inputs=["image", "image", "dropdown", "number", "number", "number", "number"], outputs="image", title="OOTDiffusion Demo DC")
+block_dc.launch(api_name='generate')