metadata
license: apache-2.0
language:
- en
pipeline_tag: text-to-image
tags:
- pytorch
- diffusers
- conditional-image-generation
- diffusion-models-class
datasets:
- dpdl-benchmark/caltech_birds2011
library_name: diffusers
# class-conditional-diffusion-cub-200
A Diffusion model on Cub 200 dataset for generating bird images.
## Usage Predict function to generate images
```python
def load_model(model_path, device):
# Initialize the same model architecture as during training
model = ClassConditionedUnet().to(device)
# Load the trained weights
model.load_state_dict(torch.load(model_path))
# Set model to evaluation mode
model.eval()
return model
def predict(model, class_label, noise_scheduler, num_samples=8, device='cuda'):
model.eval() # Ensure the model is in evaluation mode
# Prepare a batch of random noise as input
shape = (num_samples, 3, 256, 256) # Input shape: (batch_size, channels, height, width)
noisy_image = torch.randn(shape).to(device)
# Ensure class_label is a tensor and properly repeated for the batch
class_labels = torch.tensor([class_label] * num_samples, dtype=torch.long).to(device)
# Reverse the diffusion process step by step
for t in tqdm(range(49, -1, -1), desc="Reverse Diffusion Steps"): # Iterate backwards through timesteps
t_tensor = torch.tensor([t], dtype=torch.long).to(device) # Single time step for the batch
# Predict noise with the model and remove it from the image
with torch.no_grad():
noise_pred = model(noisy_image, t_tensor.expand(num_samples), class_labels) # Class conditioning here
# Step with the scheduler (model_output, timestep, sample)
noisy_image = noise_scheduler.step(noise_pred, t, noisy_image).prev_sample
# Post-process the output to get image values between [0, 1]
generated_images = (noisy_image + 1) / 2 # Rescale from [-1, 1] to [0, 1]
return generated_images
def display_images(images, num_rows=2):
# Create a grid of images
grid = torchvision.utils.make_grid(images, nrow=num_rows)
np_grid = grid.permute(1, 2, 0).cpu().numpy() # Convert to (H, W, C) format for visualization
# Plot the images
plt.figure(figsize=(12, 6))
plt.imshow(np.clip(np_grid, 0, 1)) # Clip values to ensure valid range
plt.axis('off')
plt.show()
```
Example of loading a model and generating predictions
```python
model_path = "model_epoch_0.pth" # Path to your saved model
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = load_model(model_path, device)
noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')
class_label = 1 # Example class label, change to your desired class
generated_images = predict(model, class_label, noise_scheduler, num_samples=2, device=device)
display_images(generated_images)
```