feat: check whether generating seq pic

2023-03-14 22:46:53 +08:00 · 2023-03-14 22:46:53 +08:00 · 37adcc7e97
commit 37adcc7e97
parent d3eff8e425
1 changed files with 10 additions and 8 deletions
--- a/ddpm.py
+++ b/ddpm.py
@ -53,7 +53,7 @@ class DDPM(nn.Module):

        return mu + sigma * epsilon, epsilon                                                                # (b, c, w, h)
    
-    def sample(self, model, n):
+    def sample(self, model, generate_iteration_pic=False, n=self.batch_size):
        '''
            Inputs:
                model (nn.Module): Unet instance
@ -84,6 +84,8 @@ class DDPM(nn.Module):

                x_t = first_term * (x_t-(second_term * predict_noise)) - z * beta

+                # generate 10 pic on the different denoising times
+                if generate_iteration_pic:
                    if i % (self.iteration/10) == 0:
                        p = x_t[0].cpu()
                        p = ( p.clamp(-1, 1) + 1 ) / 2