VAE mod.functions Assert #33
Description
Hi, MLC Team! I'm trying to reproduce your model with the tvm unity and newest llvm branch. When I run the build.py script like this:
python build.py --target cuda
It turned out to trigger the assert assert len(mod.functions) == 1
When I print the len(mod.functions), I find that the len of it is 8, but we only reserve the mod["subgraph_0"], I was confused that why should we set this assert?
def vae_to_image(pipe) -> tvm.IRModule:
class VAEModelWrapper(torch.nn.Module):
def __init__(self, vae):
super().__init__()
self.vae = vae
def forward(self, latents):
latents = 1 / 0.18215 * latents
z = self.vae.post_quant_conv(latents)
image = self.vae.decoder(z)
image = (image / 2 + 0.5).clamp(min=0, max=1)
image = (image.permute(0, 2, 3, 1) * 255).round()
return image
vae = pipe.vae
vae_to_image = VAEModelWrapper(vae)
z = torch.rand((1, 4, 64, 64), dtype=torch.float32)
mod = dynamo_capture_subgraphs(
vae_to_image.forward,
z,
keep_params_as_input=True,
)
assert len(mod.functions) == 1
return tvm.IRModule({"vae": mod["subgraph_0"]})And after I commented this assert and finally saved the params, I used python3 deploy.py --prompt "A photo of an astronaut riding a horse on mars." --device-name cuda to use this model. It turned out that a runtime error had triggered:
RuntimeError: Check failed: input_shape[i] == reg (4 vs. 512) : ErrorContext(fn=image_to_rgba, loc=param[0], param=x, annotation=R.Tensor((1, 512, 512, 3), dtype="float32")) match_cast error, shape[1] mismatch to specified constant.Maybe it is because of the assert issues I had ignored before?