Update README.md

README.md

@@ -20,7 +20,7 @@ library_name: diffusers
     ## Usage
     ```python
 
-# Load the model
+# Predict function to generate images
 def load_model(model_path, device):
     # Initialize the same model architecture as during training
     model = ClassConditionedUnet().to(device)
@@ -33,7 +33,7 @@ def load_model(model_path, device):
     
     return model
 
-# Predict function to generate images
+
 def predict(model, class_label, noise_scheduler, num_samples=8, device='cuda'):
     model.eval()  # Ensure the model is in evaluation mode
     
@@ -60,7 +60,6 @@ def predict(model, class_label, noise_scheduler, num_samples=8, device='cuda'):
     
     return generated_images
 
-# Display predicted images
 def display_images(images, num_rows=2):
     # Create a grid of images
     grid = torchvision.utils.make_grid(images, nrow=num_rows)
@@ -75,17 +74,9 @@ def display_images(images, num_rows=2):
 # Example of loading a model and generating predictions
 model_path = "model_epoch_0.pth"  # Path to your saved model
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-# Load the model
 model = load_model(model_path, device)
-
-# Create a noise scheduler
 noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')
-
-# Predict and generate samples for a specific class label
 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 the generated images
 display_images(generated_images)
     ```