small fixes, add comfy pbar for model loading

This commit is contained in:
kijai
2025-07-06 23:55:37 +03:00
parent a35142b006
commit d903ea7d3d
4 changed files with 31 additions and 15 deletions
+8 -4
View File
@@ -189,6 +189,8 @@ class WanControlNet(ModelMixin):
self.in_channels = controlnet_cfg["in_channels"]
self.dim = controlnet_cfg["dim"]
self.num_heads = controlnet_cfg["num_heads"]
self.quantized = controlnet_cfg["quantized"]
self.base_dtype = controlnet_cfg["base_dtype"]
if controlnet_cfg["conv_out_dim"] != controlnet_cfg["dim"]:
self.proj_in = nn.Linear(controlnet_cfg["conv_out_dim"], controlnet_cfg["dim"])
@@ -230,11 +232,13 @@ class WanControlNet(ModelMixin):
self.controlnet_mask_embedding = MaskCamEmbed(controlnet_cfg)
def forward(self, render_latent, render_mask, camera_embedding, temb, device):
def forward(self, render_latent, render_mask, camera_embedding, temb, device):
controlnet_rotary_emb = self.controlnet_rope(render_latent)
controlnet_inputs = self.controlnet_patch_embedding(render_latent.to(torch.float32)).to(render_latent.dtype)
controlnet_inputs = controlnet_inputs.to(render_latent.dtype)
controlnet_inputs = self.controlnet_patch_embedding(render_latent.to(torch.float32))
if not self.quantized:
controlnet_inputs = controlnet_inputs.to(render_latent.dtype)
else:
controlnet_inputs = controlnet_inputs.to(self.base_dtype)
controlnet_inputs = controlnet_inputs.flatten(2).transpose(1, 2)