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 einops import rearrange
from comfy import model_management 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 from comfy.ldm.modules.attention import optimized_attention
if model_management.xformers_enabled(): if model_management.xformers_enabled():
@ -276,7 +276,7 @@ class SizeEmbedder(TimestepEmbedder):
assert s.shape[0] == bs assert s.shape[0] == bs
b, dims = s.shape[0], s.shape[1] b, dims = s.shape[0], s.shape[1]
s = rearrange(s, "b d -> (b d)") 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 = 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) s_emb = rearrange(s_emb, "(b d) d2 -> b (d d2)", b=b, d=dims, d2=self.outdim)
return s_emb return s_emb

View File

@ -120,7 +120,6 @@ class PixArtMS(PixArt):
self.micro_conditioning = micro_condition self.micro_conditioning = micro_condition
if self.micro_conditioning: if self.micro_conditioning:
self.csize_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations) 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) self.ar_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
@ -223,7 +222,7 @@ class PixArtMS(PixArt):
context = context.unsqueeze(1) context = context.unsqueeze(1)
## run original forward pass ## 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 ## only return EPS
if self.pred_sigma: if self.pred_sigma:

View File

@ -728,7 +728,7 @@ class PixArt(BaseModel):
height = kwargs.get("height", None) height = kwargs.get("height", None)
if width is not None and height is not None: if width is not None and height is not None:
out["c_size"] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width]])) 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 return out