license: mit
metrics:
- mse
library_name: diffusers
tags:
- diffusion
Generate some cute (cursed) kitty cats!
import torch from diffusers import UNet2DModel, DDPMPipeline, DDPMScheduler import safetensors
model_path = "ddpm-cats/model_epoch_9.safetensors" state_dict = safetensors.torch.load_file(model_path) image_size = 128
model = UNet2DModel( sample_size=image_size, in_channels=3, out_channels=3, layers_per_block=2, block_out_channels=(128, 128, 256, 256, 512, 512), down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D"), up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"), ) model.load_state_dict(state_dict) model.eval()
scheduler = DDPMScheduler( num_train_timesteps=1000, beta_start=0.0001, beta_end=0.02, beta_schedule="linear", )
pipeline = DDPMPipeline( unet=model, scheduler=scheduler, )
image = pipeline(batch_size=1, generator=torch.manual_seed(0)).images[0].save("generated_image_128.png")