File size: 2,930 Bytes
1156de8 4a95232 1156de8 4a95232 1156de8 4a95232 1156de8 4a95232 1156de8 4a95232 1156de8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | # Create a conditional pipeline that matches the original closely
from diffusers import DDPMPipeline
import torch
class ConditionalDDPMPipeline(DDPMPipeline):
"""DDPM Pipeline with class conditioning support"""
def __call__(
self,
batch_size=1,
generator=None,
num_inference_steps=1000,
output_type="pil",
class_labels=None,
guidance_scale=1.5,
return_dict=True,
mem=False,
):
# Initialize with random noise (exactly like original)
image = torch.randn(
(
batch_size,
self.unet.config.in_channels,
self.unet.config.sample_size,
self.unet.config.sample_size,
),
generator=generator,
device=self.device,
)
# Setup the scheduler (exactly like original)
self.scheduler.set_timesteps(num_inference_steps)
if mem:
TCNP = torch.empty((batch_size, num_inference_steps), device=self.device)
# Denoising process
for i, t in enumerate(self.scheduler.timesteps):
# Only difference is we pass class_labels to the model
with torch.no_grad():
if guidance_scale > 1.0 and class_labels is not None:
# Conditional pass
cond_output = self.unet(
image, t, class_labels=class_labels
).sample
# Unconditional pass
uncond_labels = torch.zeros_like(class_labels)
uncond_output = self.unet(
image, t, class_labels=uncond_labels
).sample
# Combine with guidance scale
model_output = uncond_output + guidance_scale * (
cond_output - uncond_output
)
if mem:
# print(cond_output.squeeze().shape, uncond_output.shape)
TCNP[:, i] = torch.linalg.norm(cond_output.squeeze() - uncond_output.squeeze(), dim=[0,1])
else:
# Standard pass with conditioning
model_output = self.unet(
image, t, class_labels=class_labels
).sample
# Scheduler step (exactly like original)
image = self.scheduler.step(
model_output, t, image, generator=generator
).prev_sample
# Final processing (exactly like original)
image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
if output_type == "pil":
image = self.numpy_to_pil(image)
if not return_dict:
return (image,)
if mem:
return dict(images=image, TCNP=TCNP)
return dict(images=image)
|