update

app.py

@@ -11,9 +11,6 @@
 # // 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
@@ -24,6 +21,7 @@ import datetime
 from tqdm import tqdm
 import gc
 
+from data.image.transforms.divisible_crop import DivisibleCrop
 from data.image.transforms.na_resize import NaResize
 from data.video.transforms.rearrange import Rearrange
 if os.path.exists("./projects/video_diffusion_sr/color_fix.py"):
@@ -36,6 +34,8 @@ from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 import argparse
 
+import spaces
+
 from common.distributed import (
     get_device,
     init_torch,
@@ -60,27 +60,6 @@ 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.
 
@@ -149,7 +128,7 @@ def configure_sequence_parallel(sp_size):
     if sp_size > 1:
         init_sequence_parallel(sp_size)
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=120)
 def configure_runner(sp_size):
     config_path = os.path.join('./configs_3b', 'main.yaml')
     config = load_config(config_path)
@@ -165,7 +144,7 @@ def configure_runner(sp_size):
         runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
     return runner
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=120)
 def generation_step(runner, text_embeds_dict, cond_latents):
     def _move_to_cuda(x):
         return [i.to(torch.device("cuda")) for i in x]
@@ -222,7 +201,7 @@ def generation_step(runner, text_embeds_dict, cond_latents):
 
     return samples
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=120)
 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'