fix device and order

Author: strint

src/diffusers/oneflow_graph_compile_cache.py

@@ -64,11 +64,11 @@ def __call__(self, *args, **kwargs):
 
 
 class LRUCache(object):
+    _cnt: int = 0
     def __init__(self, cache_size):
         self.cache_size = cache_size
         self.queue = deque()
         self.hash_map = dict()
-        self.cnt = 0
 
     def front(self):
         if self.is_empty():
@@ -98,8 +98,8 @@ def set(self, key, value):
             pop_key = self.pop()
 
         self.queue.appendleft(key)
-        value._oneflow_graph_cache_order = self.cnt
-        self.cnt += 1
+        value._oneflow_graph_cache_order = LRUCache._cnt
+        LRUCache._cnt += 1
         self.hash_map[key] = value
         return pop_key if pop_key is not None else key
 
@@ -157,11 +157,13 @@ def load_graph(self, path):
             sub_folders = [ f.path for f in os.scandir(path) if f.is_dir() ]
             graph_dict = dict()
             for sub_folder in sub_folders:
-                state_dict = flow.load(sub_folder, map_location="cuda")
+                state_dict = flow.load(sub_folder)
                 cache_order = state_dict["cache_order"]
+                print("===> order", cache_order)
                 graph_dict[cache_order] = state_dict
 
             for order, state_dict in sorted(graph_dict.items()):
+                print("===> load order", order)
                 graph_class_name  = state_dict["graph_class_name"]
                 cache_key = state_dict["cache_key"]
                 if graph_class_name not in self.cache_bucket_:

src/diffusers/utils/__init__.py

@@ -14,6 +14,7 @@
 
 
 import os
+import time
 
 from .deprecation_utils import deprecate
 from .import_utils import (
@@ -87,3 +88,19 @@
     "EulerAncestralDiscreteScheduler",
     "DPMSolverMultistepScheduler",
 ]
+
+def cost_cnt(fn):
+    import oneflow as flow
+    def new_fn(*args, **kwargs):
+        print("==>", fn.__name__, " try to run")
+        before_used = flow._oneflow_internal.GetCUDAMemoryUsed()
+        start_time = time.time()
+        out = fn(*args, **kwargs)
+        end_time = time.time()
+        after_used = flow._oneflow_internal.GetCUDAMemoryUsed()
+        print(fn.__name__, " run time ", end_time - start_time)
+        print(fn.__name__, " cuda mem", after_used - before_used)
+        print("<==", fn.__name__, " finish run")
+        return out
+
+    return new_fn