diff --git a/README.md b/README.md index a32528e..f46d086 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ ## :star2: Changelog +- **[2024.4.22]** Fix unstable quality of image while multi-batch. - **[2024.4.19]** Documenting nodes - **[2024.4.19]** Initial repo diff --git a/assets/multi-batch.png b/assets/multi-batch.png new file mode 100644 index 0000000..48fd868 Binary files /dev/null and b/assets/multi-batch.png differ diff --git a/ella.py b/ella.py index 0843c7c..1707e5c 100644 --- a/ella.py +++ b/ella.py @@ -1,8 +1,10 @@ import os +from typing import Dict import folder_paths import torch from comfy import model_management +from comfy.conds import CONDCrossAttn from safetensors.torch import load_model from .model import ELLA, T5TextEmbedder @@ -37,7 +39,8 @@ class EllaProxyUNet: self.ella.to(self.dtype) for i in range(len(self.embeds)): for k in self.embeds[i]: - self.embeds[i][k].to(device=self.load_device, dtype=self.dtype) + self.embeds[i][k].to(dtype=self.dtype) + self.embeds[i][k] = CONDCrossAttn(self.embeds[i][k]) @property def load_device(self): @@ -47,10 +50,13 @@ class EllaProxyUNet: def offload_device(self): return model_management.text_encoder_offload_device() + def process_cond(self, embeds: Dict[str, CONDCrossAttn], batch_size, **kwargs): + return {k: v.process_cond(batch_size, self.load_device, **kwargs).cond for k, v in embeds.items()} + def prepare_conds(self): self.ella.to(self.load_device) - cond = self.ella(torch.Tensor([999]).to(torch.int64), **self.embeds[0]) - uncond = self.ella(torch.Tensor([999]).to(torch.int64), **self.embeds[1]) + cond = self.ella(torch.Tensor([999]).to(torch.int64), **self.process_cond(self.embeds[0], 1)) + uncond = self.ella(torch.Tensor([999]).to(torch.int64), **self.process_cond(self.embeds[1], 1)) self.ella.to(self.offload_device) return cond, uncond @@ -65,8 +71,8 @@ class EllaProxyUNet: self.ella.to(device=self.load_device) for i in cond_or_uncond: h = self.ella( - self.model_sampling.timestep(timestep_[i]), - **self.embeds[i], + self.model_sampling.timestep(timestep_[0]), + **self.process_cond(self.embeds[i], input_x.size(0) // len(cond_or_uncond)), ) time_aware_encoder_hidden_states.append(h) self.ella.to(self.offload_device) diff --git a/requirements.txt b/requirements.txt index c78a037..36ad8bb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ torch safetensors transformers +sentencepiece