rm dynamo (#558)

siliconflow · Jan 23, 2024 · 0dd6c37 · 0dd6c37
1 parent cffd3d4
commit 0dd6c37
Show file tree

Hide file tree

Showing 5 changed files with 1 addition and 11 deletions.
diff --git a/examples/text_to_image_deep_cache_sdxl_enterprise.py b/examples/text_to_image_deep_cache_sdxl_enterprise.py
@@ -4,7 +4,6 @@
 
 import torch
 import torch.nn as nn
-from torch._dynamo import allow_in_graph as maybe_allow_in_graph
 
 # oneflow_compile should be imported before importing any diffusers
 from onediff.infer_compiler import oneflow_compile
@@ -112,7 +111,6 @@ def parse_args():
         False,
         False,
         args.bits,
-        maybe_allow_in_graph,
     )
 
 if args.complie_text_encoder:

diff --git a/examples/text_to_image_sdxl_enterprise.py b/examples/text_to_image_sdxl_enterprise.py
@@ -4,7 +4,6 @@
 
 import torch
 import torch.nn as nn
-from torch._dynamo import allow_in_graph as maybe_allow_in_graph
 
 # oneflow_compile should be imported before importing any diffusers
 from onediff.infer_compiler import oneflow_compile
@@ -110,7 +109,6 @@ def parse_args():
         False,
         False,
         args.bits,
-        maybe_allow_in_graph,
     )
 
 if args.complie_text_encoder:

diff --git a/onediff_comfy_nodes/utils/onediff_quant_utils.py b/onediff_comfy_nodes/utils/onediff_quant_utils.py
@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 
-from torch._dynamo import allow_in_graph as maybe_allow_in_graph
 
 if hasattr(comfy.ops, "disable_weight_init"):
     comfy_ops_Linear = comfy.ops.disable_weight_init.Linear
@@ -241,7 +240,6 @@ def replace_module_with_quantizable_module(
             fake_quant=False,
             static=False,
             nbits=8,
-            convert_fn=maybe_allow_in_graph,
         )
         modify_sub_module(diffusion_model, sub_module_name, sub_mod)
     if use_rewrite_attn:

diff --git a/src/infer_compiler_registry/register_onediff_quant/quant_diffusion_pipeline.py b/src/infer_compiler_registry/register_onediff_quant/quant_diffusion_pipeline.py
@@ -7,7 +7,6 @@
     rewrite_sdxl_pipeline_attention,
     replace_sub_module_with_quantizable_module,
 )
-from torch._dynamo import allow_in_graph as maybe_allow_in_graph
 
 
 def _use_graph():
@@ -108,7 +107,6 @@ def _quantize_model(self):
                 self._fake_quant,
                 self._static,
                 self._bits,
-                maybe_allow_in_graph,
             )
         rewrite_sdxl_pipeline_attention(self._pipe)
 

diff --git a/src/onediff/optimization/quant_optimizer.py b/src/onediff/optimization/quant_optimizer.py
@@ -38,7 +38,6 @@ def quantize_model(
     if varify_can_use_quantization() is False:
         return model
 
-    from torch._dynamo import allow_in_graph as maybe_allow_in_graph
     from onediff_quant.utils import symm_quantize_sub_module, find_quantizable_modules
     from onediff_quant.utils import get_quantize_module
     from onediff_quant import Quantizer
@@ -90,7 +89,6 @@ def apply_quantization_to_modules(quantizable_modules):
                 fake_quant=False,
                 static=False,
                 nbits=bits,
-                convert_fn=maybe_allow_in_graph,
             )
 
             modify_sub_module(model, sub_module_name, sub_mod)
@@ -112,4 +110,4 @@ def apply_quantization_to_modules(quantizable_modules):
         + f"Time: {time.time() - start_time:.4f}s \n"
     )
 
-    return model
+    return model