Fix LCM Sampler

2024-06-28 15:52:54 +02:00
parent 92080ae196
commit bfe22d8d06
1 changed files with 2 additions and 2 deletions
@@ -110,7 +110,7 @@ class CoreMLModelWrapperLCM(CoreMLModelWrapper):
        model_input_kwargs = self.prepare_inputs(x, t, c_crossattn, control)

        np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"]
-        return (torch.from_numpy(np_out).to(x.device),)
+        return torch.from_numpy(np_out).to(x.device)

    def __call__(self, latents, ts, encoder_hidden_states, **kwargs):
-        return self.apply_model(latents, ts, c_crossattn=encoder_hidden_states)
+        return (self.apply_model(latents, ts, c_crossattn=encoder_hidden_states),)