Update gradio_tabs/animation.py

gradio_tabs/animation.py

@@ -7,17 +7,15 @@ import numpy as np
 import imageio
 import spaces
 
-extensions_dir = "./torch_extension/"
-os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
+# extensions_dir = "./torch_extension/"
+# os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
 
-from networks.generator import Generator
+# from networks.generator import Generator
 
-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.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/YaohuiW/LIA-X/resolve/main/lia-x.pt"))
-gen.eval()
+# device = torch.device("cuda")
+# gen = Generator(size=512, motion_dim=40, scale=2).to(device)
+# gen.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/YaohuiW/LIA-X/resolve/main/lia-x.pt"))
+# gen.eval()
 
 output_dir = "./res_gradio"
 os.makedirs(output_dir, exist_ok=True)
@@ -124,45 +122,41 @@ def vid_postprocessing(video, fps, output_path=output_dir + "/output_vid.mp4"):
 
 	return output_path
 
-@torch.no_grad()
-def edit_media(image, *selected_s):
 
-	image_tensor = img_preprocessing(image, 512)
-	image_tensor = image_tensor.to(device)
-
-	edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
-
-	# de-norm
-	edited_image = img_postprocessing(edited_image_tensor)
-
-	return edited_image
+def animation(gen):
+    
+	@torch.no_grad()
+	def edit_media(image, *selected_s):
 
-@torch.no_grad()
-def animate_media(image, video, *selected_s):
+		image_tensor = img_preprocessing(image, 512)
+		image_tensor = image_tensor.to(device)
 
-	image_tensor = img_preprocessing(image, 512)
-	vid_target_tensor, fps = vid_preprocessing(video, 512)
-	image_tensor = image_tensor.to(device)
-	video_target_tensor = vid_target_tensor.to(device)
+		edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
 
-	animated_video = gen.animate(image_tensor, video_target_tensor, labels_v, selected_s)
+		# de-norm
+		edited_image = img_postprocessing(edited_image_tensor)
 
-	# postprocessing
-	animated_video = vid_postprocessing(animated_video, fps)
+		return edited_image
 
-	return animated_video
+	@torch.no_grad()
+	def animate_media(image, video, *selected_s):
 
+		image_tensor = img_preprocessing(image, 512)
+		vid_target_tensor, fps = vid_preprocessing(video, 512)
+		image_tensor = image_tensor.to(device)
+		video_target_tensor = vid_target_tensor.to(device)
 
-def clear_media():
-	return None, None, *([0] * len(labels_k))
+		animated_video = gen.animate(image_tensor, video_target_tensor, labels_v, selected_s)
 
+		# postprocessing
+		animated_video = vid_postprocessing(animated_video, fps)
 
-# image_output = gr.Image(label="Output Image", elem_id="output_img", type='numpy', interactive=False, width=512)
-# video_output = gr.Video(label="Output Video", elem_id="output_vid", width=512)
+		return animated_video
 
 
+	def clear_media():
+		return None, None, *([0] * len(labels_k))
 
-def animation():
     
 	with gr.Tab("Animation & Image Editing"):