Cutouts function

app.py

@@ -23,6 +23,27 @@ from huggingface_hub import hf_hub_download
 from CLIP import clip
 from diffusion import get_model, sampling, utils
 
+class MakeCutouts(nn.Module):
+    def __init__(self, cut_size, cutn, cut_pow=1.):
+        super().__init__()
+        self.cut_size = cut_size
+        self.cutn = cutn
+        self.cut_pow = cut_pow
+
+    def forward(self, input):
+        sideY, sideX = input.shape[2:4]
+        max_size = min(sideX, sideY)
+        min_size = min(sideX, sideY, self.cut_size)
+        cutouts = []
+        for _ in range(self.cutn):
+            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)
+            offsetx = torch.randint(0, sideX - size + 1, ())
+            offsety = torch.randint(0, sideY - size + 1, ())
+            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]
+            cutout = F.adaptive_avg_pool2d(cutout, self.cut_size)
+            cutouts.append(cutout)
+        return torch.cat(cutouts)
+
 cc12m_model = hf_hub_download(repo_id="multimodalart/crowsonkb-v-diffusion-cc12m-1-cfg", filename="cc12m_1_cfg.pth")
 model = get_model('cc12m_1_cfg')()
 _, side_y, side_x = model.shape