Huertas97 commited on
Commit
55a51c2
1 Parent(s): 0ad4fb8

Upload README.md with huggingface_hub

Browse files
Files changed (1) hide show
  1. README.md +116 -0
README.md CHANGED
@@ -14,3 +14,119 @@ The code for this model can be found in this [GitHub repository](https://github.
14
 
15
  ## Usage
16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
 
15
  ## Usage
16
 
17
+ As it is a Custom Class Model of the Diffusers library, it can be used as follows:
18
+
19
+ Setup
20
+ ```python
21
+ import json
22
+ import torch
23
+ import torchvision
24
+ from matplotlib import pyplot as plt
25
+ from tqdm.auto import tqdm
26
+ from torch import nn
27
+ from diffusers import UNet2DModel, DDPMScheduler
28
+ import safetensors
29
+ from huggingface_hub import hf_hub_download
30
+ device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
31
+ ```
32
+ Load the ClassConditionedUnet model safetensor:
33
+ ```python
34
+ # Custom Class
35
+ class ClassConditionedUnet(nn.Module):
36
+ def __init__(self, num_classes=10, class_emb_size=10):
37
+ super().__init__()
38
+
39
+ # The embedding layer will map the class label to a vector of size class_emb_size
40
+ self.class_emb = nn.Embedding(num_classes, class_emb_size)
41
+
42
+ # Self.model is an unconditional UNet with extra input channels
43
+ # to accept the conditioning information (the class embedding)
44
+ self.model = UNet2DModel(
45
+ sample_size=28, # output image resolution. Equal to input resolution
46
+ in_channels=1 + class_emb_size, # Additional input channels for class cond
47
+ out_channels=1, # the number of output channels. Equal to input
48
+ layers_per_block=3, # three residual connections (ResNet) per block
49
+ block_out_channels=(128, 256, 512), # N of output channels for each block. Inverse for upsampling
50
+ down_block_types=(
51
+ "DownBlock2D", # a regular ResNet downsampling block
52
+ "AttnDownBlock2D",
53
+ "AttnDownBlock2D", # a ResNet downsampling block with spatial self-attention
54
+ ),
55
+ up_block_types=(
56
+ "AttnUpBlock2D", # a ResNet upsampling block with spatial self-attention
57
+ "AttnUpBlock2D",
58
+ "UpBlock2D", # a regular ResNet upsampling block
59
+ ),
60
+ dropout = 0.1, # Dropout prob between Conv1 and Conv2 in a block. From Improved DDPM paper
61
+ )
62
+
63
+ # Forward method takes the class labels as an additional argument
64
+ def forward(self, x, t, class_labels):
65
+ bs, ch, w, h = x.shape # x is shape (bs, 1, 28, 28)
66
+
67
+ # class conditioning embedding to add as additional input channels
68
+ class_cond = self.class_emb(class_labels) # Map to embedding dimension
69
+ class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w, h)
70
+ # class_cond final shape (bs, 4, 28, 28)
71
+
72
+ # Model input is now x and class cond concatenated together along dimension 1
73
+ # We need provide additional information (the class label)
74
+ # to every spatial location (pixel) in the image. Not changing the original
75
+ # pixels of the images, but adding new channels.
76
+ net_input = torch.cat((x, class_cond), 1) # (bs, 5, 28, 28)
77
+
78
+ # Feed this to the UNet alongside the timestep and return the prediction
79
+ # with image output size
80
+ return self.model(net_input, t).sample # (bs, 1, 28, 28)
81
+
82
+ # Define paths to download the model and scheduler
83
+ repo_name = "Huertas97/conditioned-unet-fashion-mnist-non-ema"
84
+
85
+ # Download the safetensors model file
86
+ model_file_path = hf_hub_download(repo_id=repo_name, filename="fashion_class_cond_unet_model_best.safetensors")
87
+
88
+ # # Load the Class Conditioned UNet model state dictionary
89
+ state_dict = safetensors.torch.load_file(model_file_path)
90
+ model_classcond_native = ClassConditionedUnet()
91
+ model_classcond_native.load_state_dict(state_dict).to(device)
92
+ ```
93
+
94
+ Load the DDPMScheduler:
95
+ ```python
96
+ # Download and load the scheduler configuration file
97
+ scheduler_file_path = hf_hub_download(repo_id=repo_name, filename="scheduler_config.json")
98
+
99
+ with open(scheduler_file_path, 'r') as f:
100
+ scheduler_config = json.load(f)
101
+
102
+ noise_scheduler = DDPMScheduler.from_config(scheduler_config)
103
+ ```
104
+
105
+ Use the model to generate images:
106
+
107
+ ```python
108
+ desired_class = [7] # desired class from 0 -> 9
109
+ num_samples = 2 # num of images to generate per class
110
+
111
+ # Prepare random x to start from
112
+ x = torch.randn(num_samples*len(desired_class), 1, 28, 28).to(device)
113
+
114
+ # Prepare the desired classes
115
+ y = torch.tensor([[i]*num_samples for i in desired_class]).flatten().to(device)
116
+
117
+ model_classcond_native = model_classcond_native.to(device)
118
+
119
+ # Sampling loop
120
+ for i, t in tqdm(enumerate(noise_scheduler.timesteps)):
121
+ # Get model pred
122
+ with torch.no_grad():
123
+ residual = model_classcond_native(x, t, y)
124
+
125
+ # Update sample with step
126
+ x = noise_scheduler.step(residual, t, x).prev_sample
127
+
128
+ # Show the results
129
+ fig, ax = plt.subplots(1, 1, figsize=(12, 12))
130
+ ax.imshow(torchvision.utils.make_grid(x.detach().cpu().clip(-1, 1), nrow=8)[0], cmap='Greys')
131
+ ```
132
+