Views
No views yet
!pip install git+https://github.com/huggingface/diffusers 1from huggingface_hub import notebook_login
2notebook_login()1from diffusers import UNet2DModel, DDPMScheduler
2from diffusers.utils.torch_utils import randn_tensor
3from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
4from huggingface_hub import hf_hub_download
5import torch
6import os
7from PIL import Image
8import matplotlib.pyplot as plt
9from typing import List, Optional, Tuple, Union
10
11class DDPMPipelinenew(DiffusionPipeline):
12 def __init__(self, unet, scheduler, num_classes: int):
13 super().__init__()
14 self.register_modules(unet=unet, scheduler=scheduler)
15 self.num_classes = num_classes
16 self._device = unet.device # Ensure the pipeline knows the device
17
18 @torch.no_grad()
19 def __call__(
20 self,
21 batch_size: int = 64,
22 class_labels: Optional[torch.Tensor] = None,
23 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
24 num_inference_steps: int = 1000,
25 output_type: Optional[str] = "pil",
26 return_dict: bool = True,
27 ) -> Union[ImagePipelineOutput, Tuple]:
28
29 # Ensure class_labels is on the same device as the model
30 class_labels = class_labels.to(self._device)
31 if class_labels.ndim == 0:
32 class_labels = class_labels.unsqueeze(0).expand(batch_size)
33 else:
34 class_labels = class_labels.expand(batch_size)
35
36 # Sample gaussian noise to begin loop
37 if isinstance(self.unet.config.sample_size, int):
38 image_shape = (
39 batch_size,
40 self.unet.config.in_channels,
41 self.unet.config.sample_size,
42 self.unet.config.sample_size,
43 )
44 else:
45 image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)
46
47 image = randn_tensor(image_shape, generator=generator, device=self._device)
48
49 # Set step values
50 self.scheduler.set_timesteps(num_inference_steps)
51
52 for t in self.progress_bar(self.scheduler.timesteps):
53 # Ensure the class labels are correctly broadcast to match the input tensor shape
54 model_output = self.unet(image, t, class_labels).sample
55
56 image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
57
58 image = (image / 2 + 0.5).clamp(0, 1)
59 image = image.cpu().permute(0, 2, 3, 1).numpy()
60 if output_type == "pil":
61 image = self.numpy_to_pil(image)
62
63 if not return_dict:
64 return (image,)
65
66 return ImagePipelineOutput(images=image)
67
68 def to(self, device: torch.device):
69 self._device = device
70 self.unet.to(device)
71 return self
72
73def load_pipeline(repo_id, num_classes, device):
74 unet = UNet2DModel.from_pretrained(repo_id, subfolder="unet").to(device)
75 scheduler = DDPMScheduler.from_pretrained(repo_id, subfolder="scheduler")
76 pipeline = DDPMPipelinenew(unet=unet, scheduler=scheduler, num_classes=num_classes)
77 return pipeline.to(device) # Move the entire pipeline to the device
78
79def save_images_locally(images, save_dir, epoch, class_label):
80 os.makedirs(save_dir, exist_ok=True)
81 for i, image in enumerate(images):
82 image_path = os.path.join(save_dir, f"image_epoch{epoch}_class{class_label}_idx{i}.png")
83 image.save(image_path)
84
85def generate_images(pipeline, class_label, batch_size, num_inference_steps, save_dir, epoch):
86 generator = torch.Generator(device=pipeline._device).manual_seed(0)
87 class_labels = torch.tensor([class_label] * batch_size).to(pipeline._device)
88 images = pipeline(
89 generator=generator,
90 batch_size=batch_size,
91 num_inference_steps=num_inference_steps,
92 class_labels=class_labels,
93 output_type="pil",
94 ).images
95 save_images_locally(images, save_dir, epoch, class_label)
96 return images
97
98def create_image_grid(images, grid_size, save_path):
99 assert len(images) == grid_size ** 2, "Number of images must be equal to grid_size squared"
100 width, height = images[0].size
101 grid_img = Image.new('RGB', (grid_size * width, grid_size * height))
102
103 for i, image in enumerate(images):
104 x = i % grid_size * width
105 y = i // grid_size * height
106 grid_img.paste(image, (x, y))
107
108 grid_img.save(save_path)
109 return grid_img
110
111if __name__ == "__main__":
112 repo_id = "Ketansomewhere/cifar10_conditional_diffusion1"
113 num_classes = 10 # Adjust to your number of classes
114 batch_size = 64
115 num_inference_steps = 1000 # Can be as low as 50 for faster generation
116 save_dir = "generated_images"
117 epoch = 0
118 grid_size = 8 # 8x8 grid
119
120 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
121 pipeline = load_pipeline(repo_id, num_classes, device)
122
123 for class_label in range(num_classes):
124 images = generate_images(pipeline, class_label, batch_size, num_inference_steps, save_dir, epoch)
125
126 # Create and save the grid image
127 grid_img_path = os.path.join(save_dir, f"grid_image_class{class_label}.png")
128 grid_img = create_image_grid(images, grid_size, grid_img_path)
129
130 # Plot the grid image
131 plt.figure(figsize=(10, 10))
132 plt.imshow(grid_img)
133 plt.axis('off')
134 plt.title(f'Class {class_label}')
135 plt.savefig(os.path.join(save_dir, f"grid_image_class{class_label}.png"))
136 plt.show()