Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,60 @@
+---
+license: mit
+tags:
+- pytorch
+- diffusers
+- class-conditional-image-generation
+- diffusion-models-class
+---
+
+
+
+# This model is a diffusion model for generate number.
+
+## Usage
+
+```python
+# 定义模型
+class ClassConditionalUnet(nn.Module):
+    def __init__(self, num_classes=10, class_emb_size=4):
+        super().__init__()
+        # 将数字所属的类别映射到一个长度为class_emb_size的特征向量
+        self.class_emb = nn.Embedding(num_classes, class_emb_size)
+        # self.model就是一个不带条件的unet模型，在这里给他添加额外的输入通道，用于接收条件信息
+        self.model = UNet2DModel(
+            sample_size=28,  #生成的图像是28*28
+            in_channels=1 + class_emb_size,  #加入额外的输入通道
+            out_channels=1,  # 输入单通道黑白数字图
+            layers_per_block=2,  # 设置一个unet模块有多少个残差连接层
+            block_out_channels=(32, 64, 64),
+            down_block_types=(
+                "DownBlock2D",  #普通的ResNet下采样模块
+                "AttnDownBlock2D",  #含有spatial self-attention的下采样和模块
+                "AttnDownBlock2D",
+
+            ),
+            up_block_types=(
+                "AttnUpBlock2D",  #含有spatial self-attention的ResNet上采样模块
+                "AttnUpBlock2D",
+                "UpBlock2D",
+            ),
+        )
+
+    def forward(self, x, t, class_labels):
+        bs, ch, w, h = x.shape
+        class_cond = self.class_emb(class_labels)  # 将类别映射为向量形式
+        class_cond = class_cond.view(bs, class_cond.shape[1], 1, 1).expand(bs, class_cond.shape[1], w,
+                                                                                 h)  # 拓展张量形状
+        net_input = torch.cat((class_cond, x), dim=1)
+        return self.model(net_input, t).sample        
+    ckpt = torch.load("class_cond_unet.pth", map_location="cpu")
+
+    model = ClassConditionalUnet(
+    num_classes=ckpt["num_classes"],
+    class_emb_size=ckpt["class_emb_size"]
+)
+
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+    noise_scheduler = DDPMScheduler.from_pretrained("Dirry525/class_num_generator")
+```