Fix case where pixel space VAE could cause issues. (#13030)

This commit is contained in:
comfyanonymous 2026-03-17 17:46:22 -07:00 committed by GitHub
parent 735a0465e5
commit 68d542cc06
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -952,8 +952,8 @@ class VAE:
batch_number = max(1, batch_number)
for x in range(0, samples_in.shape[0], batch_number):
samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
out = self.process_output(self.first_stage_model.decode(samples, **vae_options).to(self.output_device).to(dtype=self.vae_output_dtype()))
samples = samples_in[x:x + batch_number].to(device=self.device, dtype=self.vae_dtype)
out = self.process_output(self.first_stage_model.decode(samples, **vae_options).to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True))
if pixel_samples is None:
pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device, dtype=self.vae_output_dtype())
pixel_samples[x:x+batch_number] = out