fix some bugs

app.py

@@ -60,7 +60,7 @@ class LeffaPredictor(object):
         src_image_path,
         ref_image_path,
         control_type,
-        ref_acceleration=True,
+        ref_acceleration=False,
         step=50,
         scale=2.5,
         seed=42,

leffa/inference.py

@@ -33,18 +33,22 @@ class LeffaInference(object):
     def __call__(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
         data = self.to_gpu(data)
 
+        ref_acceleration = kwargs.get("ref_acceleration", False)
         num_inference_steps = kwargs.get("num_inference_steps", 50)
         guidance_scale = kwargs.get("guidance_scale", 2.5)
         seed = kwargs.get("seed", 42)
+        repaint = kwargs.get("repaint", False)
         generator = torch.Generator(self.pipe.device).manual_seed(seed)
         images = self.pipe(
             src_image=data["src_image"],
             ref_image=data["ref_image"],
             mask=data["mask"],
             densepose=data["densepose"],
+            ref_acceleration=ref_acceleration,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             generator=generator,
+            repaint=repaint,
         )[0]
 
         # images = [pil_to_tensor(image) for image in images]

leffa/pipeline.py

@@ -48,7 +48,7 @@ class LeffaPipeline(object):
         ref_image,
         mask,
         densepose,
-        ref_acceleration=True,
+        ref_acceleration=False,
         num_inference_steps=50,
         do_classifier_free_guidance=True,
         guidance_scale=2.5,
@@ -182,7 +182,7 @@ class LeffaPipeline(object):
             mask = numpy_to_pil(mask)
             mask = [i.convert("RGB") for i in mask]
             gen_image = [
-                repaint(_src_image, _mask, _gen_image)
+                do_repaint(_src_image, _mask, _gen_image)
                 for _src_image, _mask, _gen_image in zip(src_image, mask, gen_image)
             ]
 
@@ -216,7 +216,7 @@ def numpy_to_pil(images):
     return pil_images
 
 
-def repaint(person, mask, result):
+def do_repaint(person, mask, result):
     _, h = result.size
     kernal_size = h // 100
     if kernal_size % 2 == 0:

preprocess/humanparsing/run_parsing.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+import os
 import sys
 import onnxruntime as ort
 PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
@@ -9,14 +10,16 @@ from parsing_api import onnx_inference
 class Parsing:
     def __init__(self, atr_path, lip_path):
         session_options = ort.SessionOptions()
+        session_options.inter_op_num_threads = os.cpu_count() // 2
+        session_options.intra_op_num_threads = os.cpu_count() // 2
         session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
         session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
         self.session = ort.InferenceSession(atr_path,
                                             sess_options=session_options, providers=['CPUExecutionProvider'])
         self.lip_session = ort.InferenceSession(lip_path,
                                                 sess_options=session_options, providers=['CPUExecutionProvider'])
-        
 
     def __call__(self, input_image):
-        parsed_image, face_mask = onnx_inference(self.session, self.lip_session, input_image)
+        parsed_image, face_mask = onnx_inference(
+            self.session, self.lip_session, input_image)
         return parsed_image, face_mask