chore: Refactor CUDA device usage in app.py

app.py

@@ -25,6 +25,9 @@ from diffusers_vdm.pipeline import LatentVideoDiffusionPipeline
 from diffusers_vdm.utils import resize_and_center_crop, save_bcthw_as_mp4
 import spaces
 
+# Disable gradients globally
+torch.set_grad_enabled(False)
+
 class ModifiedUNet(UNet2DConditionModel):
     @classmethod
     def from_config(cls, *args, **kwargs):
@@ -68,7 +71,6 @@ def find_best_bucket(h, w, options):
     return best_bucket
 
 
-@torch.inference_mode()
 def encode_cropped_prompt_77tokens(txt: str):
     cond_ids = tokenizer(txt,
                          padding="max_length",
@@ -79,7 +81,6 @@ def encode_cropped_prompt_77tokens(txt: str):
     return text_cond
 
 
-@torch.inference_mode()
 def pytorch2numpy(imgs):
     results = []
     for x in imgs:
@@ -90,7 +91,6 @@ def pytorch2numpy(imgs):
     return results
 
 
-@torch.inference_mode()
 def numpy2pytorch(imgs):
     h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
     h = h.movedim(-1, 1)
@@ -103,14 +103,12 @@ def resize_without_crop(image, target_width, target_height):
     return np.array(resized_image)
 
 
-@torch.inference_mode()
 @spaces.GPU()
 def interrogator_process(x):
     image_description = wd14tagger.default_interrogator(x)
     return image_description, image_description
 
 
-@torch.inference_mode()
 @spaces.GPU()
 def process(input_fg, prompt, input_undo_steps, image_width, image_height, seed, steps, n_prompt, cfg,
             progress=gr.Progress()):
@@ -147,7 +145,6 @@ def process(input_fg, prompt, input_undo_steps, image_width, image_height, seed,
     return pixels
 
 
-@torch.inference_mode()
 def process_video_inner(image_1, image_2, prompt, seed=123, steps=25, cfg_scale=7.5, fs=3, progress_tqdm=None):
     random.seed(seed)
     np.random.seed(seed)
@@ -198,7 +195,6 @@ def process_video_inner(image_1, image_2, prompt, seed=123, steps=25, cfg_scale=
     return video, image_1, image_2
 
 
-@torch.inference_mode()
 @spaces.GPU(duration=360)
 def process_video(keyframes, prompt, steps, cfg, fps, seed, progress=gr.Progress()):
     result_frames = []