fix pip load

Author: strint

src/diffusers/modeling_oneflow_utils.py

@@ -92,20 +92,7 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
     Reads a checkpoint file, returning properly formatted errors if they arise.
     """
     try:
-        # this is oneflow saved model, a dir
-        if os.path.isdir(checkpoint_file):
-            return torch.load(checkpoint_file, map_location="cpu")
-        elif os.path.basename(checkpoint_file) == WEIGHTS_NAME:
-            import torch as og_torch
-
-            torch_parameters = og_torch.load(checkpoint_file, map_location="cpu")
-            oneflow_parameters = dict()
-            for key, value in torch_parameters.items():
-                if value.is_cuda:
-                    raise ValueError(f"torch model is not on cpu, it is on {value.device}")
-                val = value.detach().cpu().numpy()
-                oneflow_parameters[key] = torch.from_numpy(val)
-            return oneflow_parameters
+        return torch.load(checkpoint_file, map_location="cpu")
     except Exception as e:
         try:
             with open(checkpoint_file) as f:

src/diffusers/oneflow_graph_compile_cache.py

@@ -154,10 +154,10 @@ def save_graph(self, path):
 
     def load_graph(self, path, graph_class2init_args=None):
         if self.enable_load_graph_:
-            sub_folders = [ f.path for f in os.scandir(path) if f.is_dir() ]
+            sub_files = [ f.path for f in os.scandir(path) if f.is_file() ]
             graph_dict = dict()
-            for sub_folder in sub_folders:
-                state_dict = flow.load(sub_folder)
+            for sub_file in sub_files:
+                state_dict = flow.load(sub_file)
                 cache_order = state_dict["cache_order"]
                 graph_dict[cache_order] = state_dict
 

src/diffusers/pipeline_oneflow_utils.py

@@ -675,7 +675,8 @@ def load_module(name, value):
         # 3. Load each module in the pipeline
         for name, (library_name, class_name) in init_dict.items():
             if name in ["scheduler", "unet", "vae", "text_encoder", "safety_checker"]:
-                class_name = "OneFlow" + class_name
+                if "OneFlow" not in class_name:
+                    class_name = "OneFlow" + class_name
             # 3.1 - now that JAX/Flax is an official framework of the library, we might load from Flax names
             if class_name.startswith("Flax"):
                 class_name = class_name[4:]

tests/test_pipelines_oneflow_graph_load.py

@@ -1,27 +1,39 @@
 import oneflow as torch
 import time
+import os
+import shutil
+
 from diffusers import (
     OneFlowStableDiffusionPipeline as StableDiffusionPipeline,
     OneFlowEulerDiscreteScheduler as EulerDiscreteScheduler,
 )
 from diffusers import utils
 
 model_id = "stabilityai/stable-diffusion-2"
+_graph_save_file = "./test_sd_save_graph"
+_sch_file_path = "./test_sd_sch"
+_pipe_file_path = "./test_sd_pipe"
 
-_offline_compile = False
+_online_mode = True
+_pipe_from_file = True
 
 total_start_t = time.time()
 start_t = time.time()
-# StableDiffusionPipeline 需要支持 unet 和 vae load graph， 此时无需创建 eager module
 @utils.cost_cnt
 def get_pipe():
-    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
-    sd_pipe = StableDiffusionPipeline.from_pretrained(
-        model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
-        )
+    if _pipe_from_file:
+        scheduler = EulerDiscreteScheduler.from_pretrained(_sch_file_path, subfolder="scheduler")
+        sd_pipe = StableDiffusionPipeline.from_pretrained(
+            _pipe_file_path, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
+            )
+    else:
+        scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+        sd_pipe = StableDiffusionPipeline.from_pretrained(
+            model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16
+            )
     torch._oneflow_internal.eager.Sync()
-    return sd_pipe
-pipe = get_pipe()
+    return scheduler, sd_pipe
+sch, pipe = get_pipe()
 
 @utils.cost_cnt
 def pipe_to_cuda():
@@ -37,13 +49,14 @@ def config_graph():
     torch._oneflow_internal.eager.Sync()
 config_graph()
 
-if _offline_compile:
+if not _online_mode:
     pipe.enable_save_graph()
 else:
     @utils.cost_cnt
     def load_graph():
         pipe.enable_load_graph()
-        pipe.load_graph("./test_save_load", compile_unet=True, compile_vae=False)
+        assert (os.path.exists(_graph_save_file) and os.path.isdir(_graph_save_file))
+        pipe.load_graph(_graph_save_file, compile_unet=True, compile_vae=False)
         torch._oneflow_internal.eager.Sync()
     load_graph()
 end_t = time.time()
@@ -85,9 +98,19 @@ def text_to_image(prompt, image_size, num_images_per_prompt=1, prefix=""):
 total_end_t = time.time()
 print("st init and run time ", total_end_t - total_start_t, 's.')
 
+@utils.cost_cnt
+def save_pipe_sch():
+    pipe.save_pretrained(_pipe_file_path)
+    sch.save_pretrained(_sch_file_path)
+
 @utils.cost_cnt
 def save_graph():
-    pipe.save_graph("./test_save_load")
+    if os.path.exists(_graph_save_file) and os.path.isdir(_graph_save_file):
+        shutil.rmtree(_graph_save_file)
+        os.makedirs(_graph_save_file)
+
+    pipe.save_graph(_graph_save_file)
 
-if _offline_compile:
+if not _online_mode:
+    save_pipe_sch()
     save_graph()