From fe0b3929c271e73109c5dac26bfc7c12016da6dc Mon Sep 17 00:00:00 2001
From: YangJX <7196390@qq.com>
Date: Tue, 7 Jan 2025 12:51:02 +0800
Subject: [PATCH] Add files via upload

choose from inswap_128.fp16.onnx and inswap_128.onnx
---
 __init__.py    | 17 +++++++++--------
 rope/Models.py | 26 ++++++++++++++++++++------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/__init__.py b/__init__.py
index 075924b..c548886 100644
--- a/__init__.py
+++ b/__init__.py
@@ -207,8 +207,8 @@ class RopeWrapper_LoadModels:
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "hidden": {
-                "unique_id": "UNIQUE_ID"
+            "required": {
+                "inswap_type": (["Original","fp16"], ),
             },
         }
 
@@ -218,12 +218,13 @@ class RopeWrapper_LoadModels:
     CATEGORY = "RopeWrapper"
     model=None
     vm=None
-    def run(self, unique_id):
-        if self.model is None:
-            self.model = Models.Models()
-            self.model.setModelPath(os.path.dirname(os.path.realpath(__file__))+"/")
-            self.vm = VM.VideoManager(self.model)
-        
+    def run(self, inswap_type):
+        #if self.model is None:
+        self.model = Models.Models()
+        self.model.setModelPath(os.path.dirname(os.path.realpath(__file__))+"/")
+        self.model.setInswapType(inswap_type)
+        self.vm = VM.VideoManager(self.model)
+    
         return ( self.model, self.vm )
 
 class RopeWrapper_SwapNode:
diff --git a/rope/Models.py b/rope/Models.py
index f384381..ffda4d0 100644
--- a/rope/Models.py
+++ b/rope/Models.py
@@ -27,6 +27,8 @@ class Models():
 
         self.recognition_model = []
         self.swapper_model = []
+        self.swapper_model_loaded_type = 0
+        self.swapper_model_type="fp16"
         self.swapper_model_kps = []
         self.swapper_model_swap = []
 
@@ -44,7 +46,8 @@ class Models():
     
     def setModelPath(self, path):
         self.modelPath = path
-        
+    def setInswapType(self, type):
+        self.swapper_model_type = type
     def get_gpu_memory(self):
         command = "nvidia-smi --query-gpu=memory.total --format=csv"
         memory_total_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
@@ -121,10 +124,16 @@ class Models():
 
     def calc_swapper_latent(self, source_embedding):
         if not self.swapper_model:
-            print("now begin load inswapper_128.onnx")      
-            graph = onnx.load(self.modelPath+"models/inswapper_128.onnx").graph            
-            self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
-            print("now loaded inswapper_128.onnx")    
+            if self.swapper_model_type=="fp16":
+                print("calc_swapper_latent:now begin load inswapper_128.fp16.onnx")      
+                graph = onnx.load(self.modelPath+"models/inswapper_128.fp16.onnx").graph            
+                self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
+                print("now loaded inswapper_128.fp16.onnx")    
+            else:
+                print("calc_swapper_latent:now begin load inswapper_128.onnx")      
+                graph = onnx.load(self.modelPath+"models/inswapper_128.onnx").graph            
+                self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
+                print("now loaded inswapper_128.onnx")                   
         
         n_e = source_embedding / l2norm(source_embedding)
         latent = n_e.reshape((1,-1))
@@ -138,7 +147,12 @@ class Models():
             sess_options = onnxruntime.SessionOptions()
             sess_options.enable_cpu_mem_arena = False            
             # self.swapper_model = onnxruntime.InferenceSession( self.modelPath+"models/inswapper_128_last_cubic.onnx", sess_options, providers=[('CUDAExecutionProvider', cuda_options), 'CPUExecutionProvider'])            
-            self.swapper_model = onnxruntime.InferenceSession( self.modelPath+"models/inswapper_128.onnx", providers=self.providers)
+            if self.swapper_model_type=="fp16":   
+                print("run_swapper:now begin load inswapper_128.fp16.onnx")            
+                self.swapper_model = onnxruntime.InferenceSession( self.modelPath+"models/inswapper_128.fp16.onnx", providers=self.providers)
+            else:
+                print("run_swapper:now begin load inswapper_128.onnx")     
+                self.swapper_model = onnxruntime.InferenceSession( self.modelPath+"models/inswapper_128.onnx", providers=self.providers)
 
         io_binding = self.swapper_model.io_binding()
         io_binding.bind_input(name='target', device_type='cuda', device_id=0, element_type=np.float32, shape=(1,3,128,128), buffer_ptr=image.data_ptr())