⚡️ Speed up function state_dict_prefix_replace by 127%

Here's an optimized version of your Python function. The primary changes are to minimize the creation of intermediate lists and to use dictionary comprehensions for more efficient data manipulation.



### Changes and Optimizations

1. **Avoid Unneeded List Creation:** 
   - Instead of mapping and filtering the keys in a separate step (`map` and `filter`), it is done directly in the list comprehension.
   
2. **Dictionary Comprehension**: 
   - By directly assigning `out` to `{}` or `state_dict`, it forgoes unnecessary intermediate steps in the conditional initialization.
   
3. **In-Loop Item Assignment**.
   - Keys to be replaced and corresponding operations are now handled directly within loops, reducing intermediate variable assignments.

This rewritten function should perform better, especially with large dictionaries, due to reduced overhead from list operations and more efficient key manipulation.

comfy/utils.py

@@ -118,15 +118,15 @@ def state_dict_key_replace(state_dict, keys_to_replace):
     return state_dict
 
 def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False):
-    if filter_keys:
-        out = {}
-    else:
-        out = state_dict
-    for rp in replace_prefix:
-        replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp):])), filter(lambda a: a.startswith(rp), state_dict.keys())))
-        for x in replace:
-            w = state_dict.pop(x[0])
-            out[x[1]] = w
+    out = {} if filter_keys else state_dict
+
+    for old_prefix, new_prefix in replace_prefix.items():
+        keys_to_replace = [key for key in state_dict if key.startswith(old_prefix)]
+        
+        for key in keys_to_replace:
+            new_key = new_prefix + key[len(old_prefix):]
+            out[new_key] = state_dict.pop(key)
+
     return out