update with test code

Author: zRzRzRzRzRzRzR

sat/arguments.py

@@ -36,6 +36,7 @@ def add_sampling_config_args(parser):
     group.add_argument("--input-dir", type=str, default=None)
     group.add_argument("--input-type", type=str, default="cli")
     group.add_argument("--input-file", type=str, default="input.txt")
+    group.add_argument("--sampling-image-size", type=list, default=[768, 1360])
     group.add_argument("--final-size", type=int, default=2048)
     group.add_argument("--sdedit", action="store_true")
     group.add_argument("--grid-num-rows", type=int, default=1)

sat/configs/images.jpg

@@ -185,7 +185,12 @@ def decode_first_stage(self, z):
                     kwargs = {"timesteps": len(z[n * n_samples : (n + 1) * n_samples])}
                 else:
                     kwargs = {}
-                out = self.first_stage_model.decode(z[n * n_samples : (n + 1) * n_samples], **kwargs)
+                frame = z.shape[2] * 4 - 3
+                if frame <= 9:
+                    use_cp = False
+                else:
+                    use_cp = True
+                out = self.first_stage_model.decode(z[n * n_samples : (n + 1) * n_samples], use_cp=use_cp, **kwargs)
                 all_out.append(out)
         out = torch.cat(all_out, dim=0)
         return out
@@ -218,6 +223,7 @@ def sample(
         shape: Union[None, Tuple, List] = None,
         prefix=None,
         concat_images=None,
+        ofs=None,
         **kwargs,
     ):
         randn = torch.randn(batch_size, *shape).to(torch.float32).to(self.device)
@@ -241,7 +247,7 @@ def sample(
             self.model, input, sigma, c, concat_images=concat_images, **addtional_model_inputs
         )
 
-        samples = self.sampler(denoiser, randn, cond, uc=uc, scale=scale, scale_emb=scale_emb)
+        samples = self.sampler(denoiser, randn, cond, uc=uc, scale=scale, scale_emb=scale_emb, ofs=ofs)
         samples = samples.to(self.dtype)
         return samples