oneflow_compile 之后的 model 参数 update 失败

[ccssu]

Describe the bug

oneflow_compile 之后的 model 参数 update 失败。

原始的 UNetModel：

old_model = <class 'UNetModel'>(
  (input_blocks): ModuleList(
    (0): TimestepEmbedSequential(
      (0): <class 'comfy.ops.disable_weight_init.Conv2d'>(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    ))

)

# ComfyUI/custom_nodes/ComfyUI-AnimateDiff-Evolved/animatediff/motion_module_ad.py
class AnimateDiffModel(nn.Module):
    self.down_blocks = nn.VGG16Downsample()

    def inject(self, model: "UNetModel"):
        # Inject the down_blocks into the UNetModel
        model.input_blocks.insert(0, self.down_blocks)
        return model
    def set_down_blocks(self, value):
        self.down_blocks.xxx = value

经过 AnimateDiffModel 的 inject 后：

new_model = <class 'UNetModel'>(
  (input_blocks): ModuleList(
    (0): VGG16Downsample(
    )
    (1): TimestepEmbedSequential(
      (0): <class 'comfy.ops.disable_weight_init.Conv2d'>(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    ))
)

old_model = oneflow_compile(old_model)
animate_diff_model = AnimateDiffModel()
new_model = animate_diff_model.inject(old_model)

# 推理一次 
out = new_model(input)

animate_diff_model.set_down_blocks(1)
# 再次推理
out = new_model(input) #  报错 set_down_blocks 时候 oneflow_model 并没有更新

问题

• animate_diff_model.inject 之后，导致 animate_diff_model.set_down_blocks 时候 oneflow_compile 之后里面的onediff/src/onediff/infer_compiler/with_oneflow_compile.py:24 DualModule 中 _oneflow_module 并没有更新

可复现代码：

import torch 
from torch.nn import Module
import oneflow as flow
from onediff.infer_compiler import oneflow_compile
from onediff.infer_compiler.transform import register

class Demo(Module):
    def __init__(self):
        super(Demo, self).__init__()
        self.value = 0

    def forward(self, x):
        return x * self.value
    
    def set_value(self, value):
        self.value = value

class Demo2(flow.nn.Module):
    def __init__(self):
        super(Demo2, self).__init__()
        self.value = 0

    def forward(self, x):
        return x * self.value
    
    def set_value(self, value):
        self.value = value

register(torch2oflow_class_map={Demo: Demo2})

model = Demo()
of_model = oneflow_compile(model)
input_data = torch.tensor([1.0])

of_output = of_model(input_data)
model.set_value(1.0)
pt_output = model(input_data)

print("OneFlow Output:", of_output)
print("PyTorch Output:", pt_output)
# OneFlow Output: tensor([0.])
# PyTorch Output: tensor([1.])

[ccssu]

"""onediff: 
git branch: main
git commit: 3d74fe4d
"""
import torch.nn as nn
import oneflow as flow 

unet = nn.ModuleList(
    [nn.Linear(10, 20),
    nn.Linear(20, 30)]
)

from onediff.infer_compiler.with_oneflow_compile import DualModuleList,DualModule

unet = DualModuleList(unet, flow.nn.ModuleList([None] * len(unet)))

sub_model = nn.Conv2d(10, 20, 3)

unet.insert(0, sub_model)
print(f'{isinstance(unet, DualModule)}')
print(unet[0] is sub_model) # True

# unet[0]._oneflow_model ?

[strint]

This will not be fixed because non tensor is hard to update in static computation graph.