muneebable/class-conditional-diffusion-cub-200

1---2license: apache-2.03language:4- en5pipeline_tag: text-to-image6tags:7- pytorch8- diffusers9- conditional-image-generation10- diffusion-models-class11datasets:12- dpdl-benchmark/caltech_birds201113library_name: diffusers14---15 16    # class-conditional-diffusion-cub-20017    18    A Diffusion model on Cub 200 dataset for generating bird images.19    20    ## Usage Predict function to generate images21    ```python22 23      def load_model(model_path, device):24          # Initialize the same model architecture as during training25          model = ClassConditionedUnet().to(device)26          27          # Load the trained weights28          model.load_state_dict(torch.load(model_path))29          30          # Set model to evaluation mode31          model.eval()32          33          return model34      35      36      def predict(model, class_label, noise_scheduler, num_samples=8, device='cuda'):37          model.eval()  # Ensure the model is in evaluation mode38          39          # Prepare a batch of random noise as input40          shape = (num_samples, 3, 256, 256)  # Input shape: (batch_size, channels, height, width)41          noisy_image = torch.randn(shape).to(device)42          43          # Ensure class_label is a tensor and properly repeated for the batch44          class_labels = torch.tensor([class_label] * num_samples, dtype=torch.long).to(device)45      46          # Reverse the diffusion process step by step47          for t in tqdm(range(49, -1, -1), desc="Reverse Diffusion Steps"):  # Iterate backwards through timesteps48              t_tensor = torch.tensor([t], dtype=torch.long).to(device)  # Single time step for the batch49              50              # Predict noise with the model and remove it from the image51              with torch.no_grad():52                  noise_pred = model(noisy_image, t_tensor.expand(num_samples), class_labels)  # Class conditioning here53              54              # Step with the scheduler (model_output, timestep, sample)55              noisy_image = noise_scheduler.step(noise_pred, t, noisy_image).prev_sample56          57          # Post-process the output to get image values between [0, 1]58          generated_images = (noisy_image + 1) / 2  # Rescale from [-1, 1] to [0, 1]59          60          return generated_images61      62      63      def display_images(images, num_rows=2):64          # Create a grid of images65          grid = torchvision.utils.make_grid(images, nrow=num_rows)66          np_grid = grid.permute(1, 2, 0).cpu().numpy()  # Convert to (H, W, C) format for visualization67          68          # Plot the images69          plt.figure(figsize=(12, 6))70          plt.imshow(np.clip(np_grid, 0, 1))  # Clip values to ensure valid range71          plt.axis('off')72          plt.show()73    ```74 75# Example of loading a model and generating predictions76 77    ```python78    model_path = "model_epoch_0.pth"  # Path to your saved model79    device = 'cuda' if torch.cuda.is_available() else 'cpu'80    model = load_model(model_path, device)81    noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')82    class_label = 1  # Example class label, change to your desired class83    generated_images = predict(model, class_label, noise_scheduler, num_samples=2, device=device)84    display_images(generated_images)85    ```