Update gradio_tabs/animation.py

gradio_tabs/animation.py

@@ -10,17 +10,12 @@ import spaces
 extensions_dir = "./torch_extension/"
 os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
 
-@spaces.GPU(duration=240)
-def load_model():
-	from networks.generator import Generator
+from networks.generator import Generator
 
-	device = torch.device("cuda")
-	ckpt_path = './models/lia-x.pt'
-	model = Generator(size=512, motion_dim=40, scale=2).to(device)
-	model.load_state_dict(torch.load(ckpt_path, weights_only=False))
-	return model
-
-gen = load_model()
+device = torch.device("cuda")
+ckpt_path = './models/lia-x.pt'
+gen = Generator(size=512, motion_dim=40, scale=2).to(device)
+gen.load_state_dict(torch.load(ckpt_path, weights_only=False))
 gen.eval()
 
 output_dir = "./res_gradio"
@@ -129,7 +124,6 @@ def vid_postprocessing(video, fps, output_path=output_dir + "/output_vid.mp4"):
 	return output_path
 
 
-@spaces.GPU(duration=240)
 def edit_media(image, *selected_s):
 
 	image_tensor = img_preprocessing(image, 512)
@@ -143,7 +137,6 @@ def edit_media(image, *selected_s):
 	return edited_image
 
 
-@spaces.GPU(duration=240)
 def animate_media(image, video, *selected_s):
 
 	image_tensor = img_preprocessing(image, 512)