fix temp files cache and move to ZeroGPU

app.py

@@ -1,12 +1,18 @@
+# Will be fixed soon, but meanwhile:
+import os
+if os.getenv('SPACES_ZERO_GPU') == "true":
+    os.environ['SPACES_ZERO_GPU'] = "1"
+
 import gradio as gr
 import random
 import torch
-import os
 from torch import inference_mode
-from tempfile import NamedTemporaryFile
+# from tempfile import NamedTemporaryFile
+from typing import Optional
 import numpy as np
 from models import load_model
 import utils
+import spaces
 from inversion_utils import inversion_forward_process, inversion_reverse_process
 
 
@@ -31,7 +37,7 @@ def randomize_seed_fn(seed, randomize_seed):
 
 
 def invert(ldm_stable, x0, prompt_src, num_diffusion_steps, cfg_scale_src):  # , ldm_stable):
-    ldm_stable.model.scheduler.set_timesteps(num_diffusion_steps, device=device)
+    # ldm_stable.model.scheduler.set_timesteps(num_diffusion_steps, device=device)
 
     with inference_mode():
         w0 = ldm_stable.vae_encode(x0)
@@ -67,21 +73,22 @@ def sample(ldm_stable, zs, wts, steps, prompt_tar, tstart, cfg_scale_tar):  # ,
 
     return (16000, audio.squeeze().cpu().numpy())
 
-
-def edit(cache_dir,
-         input_audio,
-         model_id: str,
-         do_inversion: bool,
-         wtszs_file: str,
-         #  wts: gr.State, zs: gr.State,
-         saved_inv_model: str,
-         source_prompt="",
-         target_prompt="",
-         steps=200,
-         cfg_scale_src=3.5,
-         cfg_scale_tar=12,
-         t_start=45,
-         randomize_seed=True):
+@spaces.GPU
+def edit(
+    # cache_dir,
+    input_audio,
+    model_id: str,
+    do_inversion: bool,
+    #  wtszs_file: str,
+    wts: Optional[torch.Tensor], zs: Optional[torch.Tensor],
+    saved_inv_model: str,
+    source_prompt="",
+    target_prompt="",
+    steps=200,
+    cfg_scale_src=3.5,
+    cfg_scale_tar=12,
+    t_start=45,
+    randomize_seed=True):
 
     print(model_id)
     if model_id == LDM2:
@@ -89,7 +96,9 @@ def edit(cache_dir,
     elif model_id == LDM2_LARGE:
         ldm_stable = ldm2_large
     else:  # MUSIC
-        ldm_stable = ldm2_music
+    ldm_stable = ldm2_music
+
+    ldm_stable.model.scheduler.set_timesteps(steps, device=device)
 
     # If the inversion was done for a different model, we need to re-run the inversion
     if not do_inversion and (saved_inv_model is None or saved_inv_model != model_id):
@@ -99,29 +108,35 @@ def edit(cache_dir,
         raise gr.Error('Input audio missing!')
     x0 = utils.load_audio(input_audio, ldm_stable.get_fn_STFT(), device=device)
 
-    if not (do_inversion or randomize_seed):
-        if not os.path.exists(wtszs_file):
-            do_inversion = True
+    # if not (do_inversion or randomize_seed):
+        # if not os.path.exists(wtszs_file):
+            # do_inversion = True
             # Too much time has passed
+    if wts is None or zs is None:
+            do_inversion = True
 
     if do_inversion or randomize_seed:  # always re-run inversion
         zs_tensor, wts_tensor = invert(ldm_stable=ldm_stable, x0=x0, prompt_src=source_prompt,
                                        num_diffusion_steps=steps,
                                        cfg_scale_src=cfg_scale_src)
-        f = NamedTemporaryFile("wb", dir=cache_dir, suffix=".pth", delete=False)
-        torch.save({'wts': wts_tensor, 'zs': zs_tensor}, f.name)
-        wtszs_file = f.name
+        # f = NamedTemporaryFile("wb", dir=cache_dir, suffix=".pth", delete=False)
+        # torch.save({'wts': wts_tensor, 'zs': zs_tensor}, f.name)
+        # wtszs_file = f.name
         # wtszs_file = gr.State(value=f.name)
         # wts = gr.State(value=wts_tensor)
+        wts = wts_tensor
+        zs = zs_tensor
         # zs = gr.State(value=zs_tensor)
         # demo.move_resource_to_block_cache(f.name)
         saved_inv_model = model_id
         do_inversion = False
     else:
-        wtszs = torch.load(wtszs_file, map_location=device)
-        # wtszs = torch.load(wtszs_file.f, map_location=device)
-        wts_tensor = wtszs['wts']
-        zs_tensor = wtszs['zs']
+    #     wtszs = torch.load(wtszs_file, map_location=device)
+    #     # wtszs = torch.load(wtszs_file.f, map_location=device)
+    #     wts_tensor = wtszs['wts']
+    #     zs_tensor = wtszs['zs']
+        wts_tensor = wts.to(device)
+        zs_tensor = zs.to(device)
 
     # make sure t_start is in the right limit
     # t_start = change_tstart_range(t_start, steps)
@@ -129,7 +144,8 @@ def edit(cache_dir,
     output = sample(ldm_stable, zs_tensor, wts_tensor, steps, prompt_tar=target_prompt,
                     tstart=int(t_start / 100 * steps), cfg_scale_tar=cfg_scale_tar)
 
-    return output, wtszs_file, saved_inv_model, do_inversion
+    return output, wts.cpu(), zs.cpu(), saved_inv_model, do_inversion
+    # return output, wtszs_file, saved_inv_model, do_inversion
 
 
 def get_example():
@@ -208,7 +224,7 @@ change <code style="display:inline; background-color: lightgrey; ">duration = mi
 
 """
 
-with gr.Blocks(css='style.css', delete_cache=(3600, 3600)) as demo:
+with gr.Blocks(css='style.css') as demo: #, delete_cache=(3600, 3600)) as demo:
     def reset_do_inversion(do_inversion_user, do_inversion):
         # do_inversion = gr.State(value=True)
         do_inversion = True
@@ -219,18 +235,18 @@ with gr.Blocks(css='style.css', delete_cache=(3600, 3600)) as demo:
     def clear_do_inversion_user(do_inversion_user):
         do_inversion_user = False
         return do_inversion_user
+    
     def post_match_do_inversion(do_inversion_user, do_inversion):
         if do_inversion_user:
             do_inversion = True
             do_inversion_user = False
         return do_inversion_user, do_inversion
         
-    
     gr.HTML(intro)
-    # wts = gr.State()
-    # zs = gr.State()
+    wts = gr.State()
+    zs = gr.State()
     wtszs = gr.State()
-    cache_dir = gr.State(demo.GRADIO_CACHE)
+    # cache_dir = gr.State(demo.GRADIO_CACHE)
     saved_inv_model = gr.State()
     # current_loaded_model = gr.State(value="cvssp/audioldm2-music")
     # ldm_stable = load_model("cvssp/audioldm2-music", device, 200)
@@ -293,13 +309,13 @@ with gr.Blocks(css='style.css', delete_cache=(3600, 3600)) as demo:
         outputs=[seed], queue=False).then(
             fn=clear_do_inversion_user, inputs=[do_inversion_user], outputs=[do_inversion_user]).then(
            fn=edit,
-           inputs=[cache_dir,
+           inputs=[#cache_dir,
                    input_audio,
                    model_id,
                    do_inversion,
                    #    current_loaded_model, ldm_stable,
-                   #    wts, zs,
-                   wtszs,
+                      wts, zs,
+                #    wtszs,
                    saved_inv_model,
                    src_prompt,
                    tar_prompt,
@@ -309,7 +325,7 @@ with gr.Blocks(css='style.css', delete_cache=(3600, 3600)) as demo:
                    t_start,
                    randomize_seed
                    ],
-           outputs=[output_audio, wtszs,
+           outputs=[output_audio, wts, zs, # wtszs,
                     saved_inv_model, do_inversion]  # , current_loaded_model, ldm_stable],
         ).then(post_match_do_inversion, inputs=[do_inversion_user, do_inversion], outputs=[do_inversion_user, do_inversion]
                ).then(lambda x: (demo.temp_file_sets.append(set([str(gr.utils.abspath(x))])) if type(x) is str else None),
@@ -332,4 +348,4 @@ with gr.Blocks(css='style.css', delete_cache=(3600, 3600)) as demo:
     )
 
     demo.queue()
-    demo.launch()
+    demo.launch(state_session_capacity=15)