[BUG] How to compile another customized network(controlnet) in DiffusionPipeline?

Describe the bug
Great work! I get a lot of speedup running the standard text2imgPipeline(30%~40%). But when I run img2imgControlnetPipeline, the speedup is small (less than 10%) because controlnet is not optimized. I add the controlnet.py like deepspeed/model_implementations/diffusers/unet.py, deepspeed/module_inject/containers/unet.py, deepspeed/module_inject/replace_policy.py, deepspeed/module_inject/containers/init.py. But it didn't work, got the following error:
Traceback (most recent call last): File "/home/yinian.lw/AIGC/script/demo.py", line 164, in <module> result_imgs, iter_times = pipe(controlnet_conditioning_image=control_image, File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context return func(*args, **kwargs) File "/home/yinian.lw/AIGC/script/Hongbao520Img2Img/stable_diffusion_controlnet_img2img.py", line 831, in __call__ down_block_res_samples, mid_block_res_sample = self.controlnet( File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/model_implementations/diffusers/controlnet.py", line 40, in forward self._create_cuda_graph(*inputs, **kwargs) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/model_implementations/diffusers/controlnet.py", line 52, in _create_cuda_graph ret = self._forward(*inputs, **kwargs) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/model_implementations/diffusers/controlnet.py", line 75, in _forward return self.controlnet(sample, timestamp, encoder_hidden_states, controlnet_cond, return_dict) File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/controlnet.py", line 461, in forward sample, res_samples = downsample_block( File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/unet_2d_blocks.py", line 837, in forward hidden_states = attn( File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/transformer_2d.py", line 265, in forward hidden_states = block( File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/ops/transformer/inference/diffusers_transformer_block.py", line 91, in forward out_attn_1 = self.attn_1(out_norm_1) File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/ops/transformer/inference/diffusers_attention.py", line 188, in forward output = DeepSpeedDiffusersAttentionFunction.apply(input, context, input_mask, self.config, self.attn_qkvw, File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/ops/transformer/inference/diffusers_attention.py", line 88, in forward output = selfAttention_fp(input, context, input_mask) File "/home/yinian.lw/AIGC/DeepSpeed/deepspeed/ops/transformer/inference/diffusers_attention.py", line 61, in selfAttention_fp qkv_out = linear_func(input, attn_qkvw, attn_qkvb if attn_qkvb is not None else attn_qkvw, attn_qkvb RuntimeError: The specified pointer resides on host memory and is not registered with any CUDA device.

Could someone teach me how to make controlnet work under deepspeed, thanks ~

Start with the added deepspeed/model_implementations/diffusers/controlnet.py and compare it with the referenced unet.py, containers/unet.py, replace_policy.py, and containers/init.py. Reproduce the img2img ControlNet pipeline failure and trace the call into diffusers_attention.py, where the host-memory pointer error occurs. Done means ControlNet runs under DeepSpeed without that CUDA error and its inference path is exercised.