fix bug of dual module (setattr, and compatible with DualModule input)

onediff_diffusers_extensions/README.md

@@ -239,7 +239,6 @@ The results are shown below
 | watercolor_v1_sdxl_lora.safetensors     | 12M  | 1.54 s           | 2.01 s                       | **0.35 s**                      |                                                                         |
 
 
-
 ### Note
 
 1. OneDiff extensions for LoRA is currently not supported for PEFT, and only supports diffusers of at least version 0.21.0.

onediff_diffusers_extensions/onediffx/lora/lora.py

@@ -172,7 +172,7 @@ def load_and_fuse_lora(
             # or add_{k,v,q,out_proj}_proj_lora layers.
             rank = value_dict["lora.down.weight"].shape[0]
 
-            if isinstance(attn_processor, LoRACompatibleConv):
+            if isinstance(attn_processor, (LoRACompatibleConv, torch.nn.Conv2d)):
                 conv_fuse_lora(
                     attn_processor,
                     value_dict,
@@ -182,7 +182,7 @@ def load_and_fuse_lora(
                     offload_device=offload_device,
                     offload_weight=offload_weight,
                 )
-            elif isinstance(attn_processor, LoRACompatibleLinear):
+            elif isinstance(attn_processor, (LoRACompatibleLinear, torch.nn.Linear)):
                 linear_fuse_lora(
                     attn_processor,
                     value_dict,

src/onediff/infer_compiler/with_oneflow_compile.py

@@ -93,9 +93,8 @@ def __getattr__(self, name):
             else getattr(self._oneflow_module, name)
         )
         if isinstance(torch_attr, torch.nn.ModuleList):
-            oneflow_attr = (
-                [None] * len(torch_attr) if oneflow_attr is None else oneflow_attr
-            )
+            if oneflow_attr is None:
+                oneflow_attr = flow.nn.ModuleList([None] * len(torch_attr))
             return DualModuleList(torch_attr, oneflow_attr)
 
         elif isinstance(torch_attr, torch.nn.Module):
@@ -159,15 +158,24 @@ def __setattr__(self, key, value):
 
 
 def get_mixed_dual_module(module_cls):
+    if issubclass(module_cls, DualModule) and "MixedDualModule" in module_cls.__name__:
+        return module_cls
+
     class MixedDualModule(DualModule, module_cls):
         def __init__(self, torch_module, oneflow_module):
+            while isinstance(torch_module, DualModule):
+                torch_module = torch_module._torch_module
             DualModule.__init__(self, torch_module, oneflow_module)
 
         def _get_name(self) -> str:
             return f"{self.__class__.__name__}(of {module_cls.__name__})"
 
     return MixedDualModule
 
+@torch2oflow.register
+def _(mod: DualModule, verbose=False):
+    return torch2oflow(mod._torch_module, verbose)
+
 
 def handle_deployable_exception(func):
     @wraps(func)