Unet模型转ncnn异常

[Baiyuetribe]

大佬，请教下Unet模型转换问题，pt文件转trace后，再使用pnnx转换异常

from diffusers import StableDiffusionPipeline
import torch
import os

model_id = "dreamlike-art/dreamlike-photoreal-2.0"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

unet = pipe.unet
unet.eval()

with torch.inference_mode(), torch.autocast("cuda"):
    inputs = (
        torch.randn(1, 4, 64, 64, dtype=torch.half, device="cuda"),
        torch.randn(1, dtype=torch.half, device="cuda"),
        torch.randn(1, 77, 768, dtype=torch.half, device="cuda"),
    )
    unet_traced = torch.jit.trace(unet, inputs, strict=False)
    unet_traced.save("unet_traced.pt")

os.system("pnnx .\unet_traced.pt inputshape=[1,4,64,64],[1]i64,[1,77,768]")

# prompt = "photo, a church in the middle of a field of crops, bright cinematic lighting, gopro, fisheye lens"
# image = pipe(prompt).images[0]

# image.save("./result.jpg")

上述代码可正常输出unet_traced.pt文件，但是pnnx转换时中途异常退出，内容如下：

pnnxparam = .\unet_traced.pnnx.param
pnnxbin = .\unet_traced.pnnx.bin
pnnxpy = .\unet_traced_pnnx.py
pnnxonnx = .\unet_traced.pnnx.onnx
ncnnparam = .\unet_traced.ncnn.param
customop =
moduleop =
############# pass_level0
inline module = diffusers.models.attention.BasicTransformerBlock
inline module = diffusers.models.attention.CrossAttention
inline module = diffusers.models.attention.FeedForward
inline module = diffusers.models.attention.GEGLU
inline module = diffusers.models.attention.Transformer2DModel
inline module = diffusers.models.embeddings.TimestepEmbedding
inline module = diffusers.models.embeddings.Timesteps
inline module = diffusers.models.resnet.Downsample2D
inline module = diffusers.models.resnet.ResnetBlock2D
inline module = diffusers.models.resnet.Upsample2D
inline module = diffusers.models.unet_2d_blocks.CrossAttnDownBlock2D
inline module = diffusers.models.unet_2d_blocks.CrossAttnUpBlock2D
inline module = diffusers.models.unet_2d_blocks.DownBlock2D
inline module = diffusers.models.unet_2d_blocks.UNetMidBlock2DCrossAttn
inline module = diffusers.models.unet_2d_blocks.UpBlock2D
inline module = diffusers.models.attention.BasicTransformerBlock
inline module = diffusers.models.attention.CrossAttention
inline module = diffusers.models.attention.FeedForward
inline module = diffusers.models.attention.GEGLU
inline module = diffusers.models.attention.Transformer2DModel
inline module = diffusers.models.embeddings.TimestepEmbedding
inline module = diffusers.models.embeddings.Timesteps
inline module = diffusers.models.resnet.Downsample2D
inline module = diffusers.models.resnet.ResnetBlock2D
inline module = diffusers.models.resnet.Upsample2D
inline module = diffusers.models.unet_2d_blocks.CrossAttnDownBlock2D
inline module = diffusers.models.unet_2d_blocks.CrossAttnUpBlock2D
inline module = diffusers.models.unet_2d_blocks.DownBlock2D
inline module = diffusers.models.unet_2d_blocks.UNetMidBlock2DCrossAttn
inline module = diffusers.models.unet_2d_blocks.UpBlock2D

----------------
terminate called after throwing an instance of 'c10::NotImplementedError'
  what():  The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/diffusers/models/embeddings.py", line 8, in forward
  def forward(self: __torch__.diffusers.models.embeddings.Timesteps,
    timesteps: Tensor) -> Tensor:
    _0 = torch.arange(0, 160, dtype=6, layout=None, device=torch.device("cuda:0"), pin_memory=False)
         ~~~~~~~~~~~~ <--- HERE
    exponent = torch.mul(_0, CONSTANTS.c0)
    exponent0 = torch.div(exponent, CONSTANTS.c1)

[ferqui]

Hi, I manage to converted it using the following code:

from diffusers import StableDiffusionPipeline
import torch
import os

model_id = "dreamlike-art/dreamlike-photoreal-2.0"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)

unet = pipe.unet
unet.eval()

inputs = (
	torch.randn(1, 4, 64, 64),
	torch.randn(1),
	torch.randn(1, 77, 768),
	torch.rand(1),
	torch.rand(1)
)

class newUnet(torch.nn.Module):
	def __init__(self, unet):
		super(newUnet, self).__init__()
		self.unet = unet

	def forward(self, in0, in1, in2, c_in, c_out):
		outout = in0 + unet(in0 * c_in, in1, in2).sample * c_out
		return outout

model = newUnet(unet)
model.eval()

unet_traced = torch.jit.trace(model, inputs, strict=False)
unet_traced.save("model.pt")

And convert it to ncnn using ./pnnx ./model.pt inputshape=[1,4,64,64],[1]i64,[1,77,768],[1],[1] fp16=1 optlevel=2

But got the following error when trying to run it

layer Tensor.expand not exists or registered
Segmentation fault

[Baiyuetribe]

Same as you now, I'm wondering if some MHA op aren't on board?
https://github.com/Birch-san/stable-diffusion/pull/4/files
or just use xformers pipe.enable_xformers_memory_efficient_attention()

[EdVince]

一般来说，你是无法导出一个跟我所提供的模型一致的模型的。第一，我使用的不是diffusers库，我是用的stable-diffusion-webui的代码的；第二，直接导出是一定会失败的，我前前后后调整了一周的代码。

考虑到导出是一件十分困难的事情，所以我并不打算放出导出的过程。但如果你们一定要自己导出模型的话，我的建议是基于diffusers进行导出，因为diffusers的代码比较简洁，遇到问题也好调整。

即便是你成功导出了模型，你的模型结构也是与我的不一样的，因为我所提供的模型里还合并了一些与denoisor相关的代码，你要么就改我提供的c++代码，要么就把模型结构改成跟我的一样。最后，你导出的模型可能会在multiheadattention上遇到问题，导致速度慢、无法支持动态shape等。

总结：导出很麻烦，不建议自行导出；自行导出成功后，我的c++代码不再适用；自行导出成功后，速度、动态shape等支持可能会遇到问题。

[EdVince]

Hi, I manage to converted it using the following code:

from diffusers import StableDiffusionPipeline
import torch
import os

model_id = "dreamlike-art/dreamlike-photoreal-2.0"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)

unet = pipe.unet
unet.eval()

inputs = (
	torch.randn(1, 4, 64, 64),
	torch.randn(1),
	torch.randn(1, 77, 768),
	torch.rand(1),
	torch.rand(1)
)

class newUnet(torch.nn.Module):
	def __init__(self, unet):
		super(newUnet, self).__init__()
		self.unet = unet

	def forward(self, in0, in1, in2, c_in, c_out):
		outout = in0 + unet(in0 * c_in, in1, in2).sample * c_out
		return outout

model = newUnet(unet)
model.eval()

unet_traced = torch.jit.trace(model, inputs, strict=False)
unet_traced.save("model.pt")

And convert it to ncnn using ./pnnx ./model.pt inputshape=[1,4,64,64],[1]i64,[1,77,768],[1],[1] fp16=1 optlevel=2

But got the following error when trying to run it

layer Tensor.expand not exists or registered
Segmentation fault

Please note that, for the current code of diffusers, you can not use it directly to export the ncnn model. You will have to modify the code of diffusers.

[Baiyuetribe]

@EdVince 那有空试下这个，目前来说非二次元里最好的模型 https://huggingface.co/dreamlike-art/dreamlike-photoreal-2.0