update

app.py

@@ -11,6 +11,9 @@
 # // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # // See the License for the specific language governing permissions and
 # // limitations under the License.
+from torchvision.transforms import functional as TVF
+import spaces
+from typing import Union
 import os
 import torch
 import mediapy
@@ -58,6 +61,27 @@ from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file, get_dir
 
+class DivisibleCrop:
+    def __init__(self, factor):
+        if not isinstance(factor, tuple):
+            factor = (factor, factor)
+
+        self.height_factor, self.width_factor = factor[0], factor[1]
+
+    def __call__(self, image: Union[torch.Tensor, Image.Image]):
+        if isinstance(image, torch.Tensor):
+            height, width = image.shape[-2:]
+        elif isinstance(image, Image.Image):
+            width, height = image.size
+        else:
+            raise NotImplementedError
+
+        cropped_height = height - (height % self.height_factor)
+        cropped_width = width - (width % self.width_factor)
+
+        image = TVF.center_crop(img=image, output_size=(cropped_height, cropped_width))
+        return image
+
 def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
     """Load file form http url, will download models if necessary.
 
@@ -126,6 +150,7 @@ def configure_sequence_parallel(sp_size):
     if sp_size > 1:
         init_sequence_parallel(sp_size)
 
+@spaces.GPU(duration=90)
 def configure_runner(sp_size):
     config_path = os.path.join('./configs_3b', 'main.yaml')
     config = load_config(config_path)
@@ -141,6 +166,7 @@ def configure_runner(sp_size):
         runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
     return runner
 
+@spaces.GPU(duration=90)
 def generation_step(runner, text_embeds_dict, cond_latents):
     def _move_to_cuda(x):
         return [i.to(torch.device("cuda")) for i in x]
@@ -197,6 +223,7 @@ def generation_step(runner, text_embeds_dict, cond_latents):
 
     return samples
 
+@spaces.GPU(duration=90)
 def generation_loop(video_path='./test_videos', output_dir='./results', seed=666, batch_size=1, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, res_h=1280, res_w=720, sp_size=1):
     runner = configure_runner(1)
     output_dir = 'output/out.mp4'