From b68b9d15305be6261f79c8006d3d2afb1e3c7e8d Mon Sep 17 00:00:00 2001 From: JettHu Date: Sat, 20 Apr 2024 08:46:42 +0800 Subject: [PATCH] fix: c_crossattn was placed on the wrong device in `fp32-text-enc + force-fp32` [issues #10](https://github.com/TencentQQGYLab/ComfyUI-ELLA/issues/10) --- ella.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ella.py b/ella.py index e6abeed..0843c7c 100644 --- a/ella.py +++ b/ella.py @@ -59,6 +59,7 @@ class EllaProxyUNet: timestep_ = kwargs["timestep"] c = kwargs["c"] cond_or_uncond = kwargs["cond_or_uncond"] # [0|1] + _device = c["c_crossattn"].device time_aware_encoder_hidden_states = [] self.ella.to(device=self.load_device) @@ -70,7 +71,7 @@ class EllaProxyUNet: time_aware_encoder_hidden_states.append(h) self.ella.to(self.offload_device) - c["c_crossattn"] = torch.cat(time_aware_encoder_hidden_states, dim=0) + c["c_crossattn"] = torch.cat(time_aware_encoder_hidden_states, dim=0).to(_device) return apply_model(input_x, timestep_, **c)