Upload README.md with huggingface_hub
Browse files
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 |
+
|