Update app.py

app.py

@@ -12,23 +12,7 @@ from ganime.model.vqgan_clean.experimental.net2net_v3 import Net2Net
 
 IMAGE_SHAPE = (64, 128, 3)
 
-
-hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.data-00000-of-00001", subfolder="vqgan_kny_image_full") 
-hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.index", subfolder="vqgan_kny_image_full") 
-vqgan_path = hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint", subfolder="vqgan_kny_image_full") 
-
-
-
-hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.data-00000-of-00001", subfolder="ganime_kny_video_full") 
-hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.index", subfolder="ganime_kny_video_full") 
-gpt_path = hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint", subfolder="ganime_kny_video_full") 
-
-cfg = omegaconf.OmegaConf.load(here("configs/kny_video_gpt2_large_gradio.yaml"))
-cfg["model"]["first_stage_config"]["checkpoint_path"] = vqgan_path
-cfg["model"]["transformer_config"]["checkpoint_path"] = gpt_path
-
-model = Net2Net(**cfg["model"], trainer_config=cfg["train"], num_replicas=1)
-model.first_stage_model.build((20, *IMAGE_SHAPE))
+model = None
 
 
 # def save_video(video):
@@ -70,8 +54,29 @@ def normalize(image):
 
     return image
 
+def load_model():
+    hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.data-00000-of-00001", subfolder="vqgan_kny_image_full") 
+    hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.index", subfolder="vqgan_kny_image_full") 
+    vqgan_path = hf_hub_download(repo_id="Kurokabe/VQGAN_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint", subfolder="vqgan_kny_image_full") 
+    
+    hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.data-00000-of-00001", subfolder="ganime_kny_video_full") 
+    hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint.index", subfolder="ganime_kny_video_full") 
+    gpt_path = hf_hub_download(repo_id="Kurokabe/GANime_Kimetsu-no-yaiba_Tensorflow", filename="checkpoint", subfolder="ganime_kny_video_full") 
+    
+    cfg = omegaconf.OmegaConf.load(here("configs/kny_video_gpt2_large_gradio.yaml"))
+    cfg["model"]["first_stage_config"]["checkpoint_path"] = vqgan_path
+    cfg["model"]["transformer_config"]["checkpoint_path"] = gpt_path
+    
+    model = Net2Net(**cfg["model"], trainer_config=cfg["train"], num_replicas=1)
+    model.first_stage_model.build((20, *IMAGE_SHAPE))
+    return model
+
 
 def generate(first, last, n_frames):
+    global model
+
+    if model is None:
+        model = load_model()
     # n_frames = 20
     n_frames = int(n_frames)
     first = resize_if_necessary(first)