Merge pull request CompVis#35 from enzymezoo-code/inpainting_1.0

Inpainting 1.0

Author: enzymezoo-code

Deforum_Stable_Diffusion.py

@@ -144,14 +144,10 @@ def forward(self, x, sigma, uncond, cond, cond_scale):
 def add_noise(sample: torch.Tensor, noise_amt: float):
     return sample + torch.randn(sample.shape, device=sample.device) * noise_amt
 
-def get_output_folder(output_path,batch_folder=None):
-    yearMonth = time.strftime('%Y-%m/')
-    out_path = os.path.join(output_path,yearMonth)
+def get_output_folder(output_path, batch_folder):
+    out_path = os.path.join(output_path,time.strftime('%Y-%m/'))
     if batch_folder != "":
-        out_path = os.path.join(out_path,batch_folder)
-        # we will also make sure the path suffix is a slash if linux and a backslash if windows
-        if out_path[-1] != os.path.sep:
-            out_path += os.path.sep
+        out_path = os.path.join(out_path, batch_folder)
     os.makedirs(out_path, exist_ok=True)
     return out_path
 
@@ -184,10 +180,9 @@ def load_mask(path, shape):
     return mask
 
 def maintain_colors(prev_img, color_match_sample, mode):
-    color_coherence = 'Match Frame 0 RGB' #@param ['None', 'Match Frame 0 HSV', 'Match Frame 0 LAB', 'Match Frame 0 RGB'] {type:'string'}
-    if color_coherence == 'Match Frame 0 RGB':
+    if mode == 'Match Frame 0 RGB':
         return match_histograms(prev_img, color_match_sample, multichannel=True)
-    elif color_coherence == 'Match Frame 0 HSV':
+    elif mode == 'Match Frame 0 HSV':
         prev_img_hsv = cv2.cvtColor(prev_img, cv2.COLOR_RGB2HSV)
         color_match_hsv = cv2.cvtColor(color_match_sample, cv2.COLOR_RGB2HSV)
         matched_hsv = match_histograms(prev_img_hsv, color_match_hsv, multichannel=True)
@@ -316,57 +311,56 @@ def generate(args, return_latent=False, return_sample=False, return_c=False):
     with torch.no_grad():
         with precision_scope("cuda"):
             with model.ema_scope():
-                for n in range(args.n_samples):
-                    for prompts in data:
-                        uc = None
-                        if args.scale != 1.0:
-                            uc = model.get_learned_conditioning(batch_size * [""])
-                        if isinstance(prompts, tuple):
-                            prompts = list(prompts)
-                        c = model.get_learned_conditioning(prompts)
-
-                        if args.init_c != None:
-                          c = args.init_c
-
-                        if args.sampler in ["klms","dpm2","dpm2_ancestral","heun","euler","euler_ancestral"]:
-                            samples = sampler_fn(
-                                c=c, 
-                                uc=uc, 
-                                args=args, 
-                                model_wrap=model_wrap, 
-                                init_latent=init_latent, 
-                                t_enc=t_enc, 
-                                device=device, 
-                                cb=callback)
+                for prompts in data:
+                    uc = None
+                    if args.scale != 1.0:
+                        uc = model.get_learned_conditioning(batch_size * [""])
+                    if isinstance(prompts, tuple):
+                        prompts = list(prompts)
+                    c = model.get_learned_conditioning(prompts)
+
+                    if args.init_c != None:
+                        c = args.init_c
+
+                    if args.sampler in ["klms","dpm2","dpm2_ancestral","heun","euler","euler_ancestral"]:
+                        samples = sampler_fn(
+                            c=c, 
+                            uc=uc, 
+                            args=args, 
+                            model_wrap=model_wrap, 
+                            init_latent=init_latent, 
+                            t_enc=t_enc, 
+                            device=device, 
+                            cb=callback)
+                    else:
+                        # args.sampler == 'plms' or args.sampler == 'ddim':
+                        if init_latent is not None and args.strength > 0:
+                            z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(device))
                         else:
-                            # args.sampler == 'plms' or args.sampler == 'ddim':
-                            if init_latent is not None and args.strength > 0:
-                                z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(device))
-                            else:
-                                z_enc = torch.randn([args.n_samples, args.C, args.H // args.f, args.W // args.f], device=device)
-                            samples = sampler.decode(z_enc, 
-                                                     c, 
-                                                     t_enc, 
-                                                     unconditional_guidance_scale=args.scale,
-                                                     unconditional_conditioning=uc,
-                                                     img_callback=callback)
-
-                        if return_latent:
-                            results.append(samples.clone())
-
-                        x_samples = model.decode_first_stage(samples)
-                        if return_sample:
-                            results.append(x_samples.clone())
-
-                        x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
-
-                        if return_c:
-                            results.append(c.clone())
-
-                        for x_sample in x_samples:
-                            x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                            image = Image.fromarray(x_sample.astype(np.uint8))
-                            results.append(image)
+                            z_enc = torch.randn([args.n_samples, args.C, args.H // args.f, args.W // args.f], device=device)
+                        samples = sampler.decode(z_enc, 
+                                                 c, 
+                                                 t_enc, 
+                                                 unconditional_guidance_scale=args.scale,
+                                                 unconditional_conditioning=uc,
+                                                 img_callback=callback)
+
+                    if return_latent:
+                        results.append(samples.clone())
+
+                    x_samples = model.decode_first_stage(samples)
+                    if return_sample:
+                        results.append(x_samples.clone())
+
+                    x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+
+                    if return_c:
+                        results.append(c.clone())
+
+                    for x_sample in x_samples:
+                        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
+                        image = Image.fromarray(x_sample.astype(np.uint8))
+                        results.append(image)
     return results
 
 def sample_from_cv2(sample: np.ndarray) -> torch.Tensor:
@@ -685,7 +679,7 @@ def DeforumArgs():
     seed_behavior = "iter" #@param ["iter","fixed","random"]
 
     #@markdown **Grid Settings**
-    make_grid = True #@param {type:"boolean"}
+    make_grid = False #@param {type:"boolean"}
     grid_rows = 2 #@param 
 
     precision = 'autocast'