Fix aspect ratio logic

This commit is contained in:
City 2024-12-14 23:06:02 +01:00
parent 51dded3398
commit 41b7be08dc
3 changed files with 4 additions and 5 deletions

View File

@ -7,7 +7,7 @@ import torch.nn.functional as F
from einops import rearrange
from comfy import model_management
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, Mlp
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, Mlp, timestep_embedding
from comfy.ldm.modules.attention import optimized_attention
if model_management.xformers_enabled():
@ -276,7 +276,7 @@ class SizeEmbedder(TimestepEmbedder):
assert s.shape[0] == bs
b, dims = s.shape[0], s.shape[1]
s = rearrange(s, "b d -> (b d)")
s_freq = self.timestep_embedding(s, self.frequency_embedding_size)
s_freq = timestep_embedding(s, self.frequency_embedding_size)
s_emb = self.mlp(s_freq.to(s.dtype))
s_emb = rearrange(s_emb, "(b d) d2 -> b (d d2)", b=b, d=dims, d2=self.outdim)
return s_emb

View File

@ -120,7 +120,6 @@ class PixArtMS(PixArt):
self.micro_conditioning = micro_condition
if self.micro_conditioning:
self.csize_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
self.ar_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
@ -223,7 +222,7 @@ class PixArtMS(PixArt):
context = context.unsqueeze(1)
## run original forward pass
out = self.forward_orig(x, timesteps, context, c_size=c_size, ar=c_ar)
out = self.forward_orig(x, timesteps, context, c_size=c_size, c_ar=c_ar)
## only return EPS
if self.pred_sigma:

View File

@ -728,7 +728,7 @@ class PixArt(BaseModel):
height = kwargs.get("height", None)
if width is not None and height is not None:
out["c_size"] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width]]))
out["ar"] = comfy.conds.CONDRegular(torch.FloatTensor([[kwargs.get("ar", height/width)]]))
out["c_ar"] = comfy.conds.CONDRegular(torch.FloatTensor([[kwargs.get("aspect_ratio", height/width)]]))
return out