From 4e4d464de09fd0e99349f9e784723bd288e91aa2 Mon Sep 17 00:00:00 2001
From: Wang Yi <53533850+marigoold@users.noreply.github.com>
Date: Fri, 17 May 2024 16:27:48 +0800
Subject: [PATCH] fix bug of onediff.infer_compiler.register (#885)

---
 onediff_sd_webui_extensions/compile_ldm.py | 4 ++--
 onediff_sd_webui_extensions/compile_sgm.py | 4 ++--
 onediff_sd_webui_extensions/compile_vae.py | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/onediff_sd_webui_extensions/compile_ldm.py b/onediff_sd_webui_extensions/compile_ldm.py
index cea7db661..e6a3aec06 100644
--- a/onediff_sd_webui_extensions/compile_ldm.py
+++ b/onediff_sd_webui_extensions/compile_ldm.py
@@ -1,7 +1,7 @@
 import os
 import oneflow as flow
-from onediff.infer_compiler import oneflow_compile, register
-from onediff.infer_compiler.transform import proxy_class
+from onediff.infer_compiler import oneflow_compile
+from onediff.infer_compiler.transform import proxy_class, register
 
 from ldm.modules.attention import (
     BasicTransformerBlock,
diff --git a/onediff_sd_webui_extensions/compile_sgm.py b/onediff_sd_webui_extensions/compile_sgm.py
index 917e47038..12398a737 100644
--- a/onediff_sd_webui_extensions/compile_sgm.py
+++ b/onediff_sd_webui_extensions/compile_sgm.py
@@ -1,6 +1,6 @@
 import oneflow as flow
-from onediff.infer_compiler import oneflow_compile, register
-from onediff.infer_compiler.transform import proxy_class
+from onediff.infer_compiler import oneflow_compile
+from onediff.infer_compiler.transform import proxy_class, register
 from sd_webui_onediff_utils import (
     CrossAttentionOflow,
     GroupNorm32Oflow,
diff --git a/onediff_sd_webui_extensions/compile_vae.py b/onediff_sd_webui_extensions/compile_vae.py
index d406a4e4c..d5c9c7f26 100644
--- a/onediff_sd_webui_extensions/compile_vae.py
+++ b/onediff_sd_webui_extensions/compile_vae.py
@@ -1,8 +1,8 @@
 from modules import shared
 from modules.sd_vae_approx import model as get_vae_model, sd_vae_approx_models
 from modules.sd_vae_approx import VAEApprox
-from onediff.infer_compiler import oneflow_compile, register
-from onediff.infer_compiler.transform import proxy_class
+from onediff.infer_compiler import oneflow_compile
+from onediff.infer_compiler.transform import proxy_class, register
 
 __all__ = ["VaeCompileCtx"]