Add DDIM sampling function for image generation

model.py

@@ -165,4 +165,27 @@ class EMA():
             if param.require_grad:
                 assert name in self.backup 
                 param.data.copy_(self.backup[name])
-                self.backup = {} 
+                self.backup = {}
+
+#DDIM 
+@torch.no_grad()
+def sample_ddim(model, T,img_size,  batch_size=16, channels=3,ddim_step=50, device='cpu'):
+    model.eval()
+    alphas_cumprod_device = alphas_cumprod.to(device)
+    timestep = torch.linspace(0, T-1, steps=ddim_step, dtype=torch.long, device=device)
+    img = torch.randn(batch_size, channels, img_size, img_size).to(device)
+    
+    for i in reversed(range(len(timestep))):
+        t_current =timestep[i]
+        t_preq = timestep[i-1] if i > 0 else -1
+        t_tensor = torch.full((batch_size,), t_current, dtype=torch.long, device=device)
+        
+        pred_noise = model(img , t_tensor)
+        alphas_bar_t = alphas_cumprod_device[t_current]
+        alphas_bar_t_preq = alphas_cumprod_device[t_preq] if t_preq >= 0 else torch.tensor(1.0, device=device)
+        pred_x0 = (img - torch.sqrt(1 - alphas_bar_t) * pred_noise)/(torch.sqrt(alphas_bar_t))
+        pred_x0 = torch.clamp(pred_x0, -1.0, 1.0)
+        pred_dir = torch.sqrt(1 - alphas_bar_t_preq) * pred_noise
+        img = torch.sqrt(alphas_bar_t_preq) * pred_x0 + pred_dir
+        
+    return img