ZeroGPU (#3)

- ZeroGPU (bc257b55dd9332eb88087037c3860874746af102)


Co-authored-by: hysts <[email protected]>

.gitattributes

@@ -33,21 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-data/driving/driving1.mp4 filter=lfs diff=lfs merge=lfs -text
-data/driving/driving2.mp4 filter=lfs diff=lfs merge=lfs -text
-data/driving/driving3.mp4 filter=lfs diff=lfs merge=lfs -text
-data/driving/driving4.mp4 filter=lfs diff=lfs merge=lfs -text
-data/driving/driving9.mp4 filter=lfs diff=lfs merge=lfs -text
-data/source/einstein.png filter=lfs diff=lfs merge=lfs -text
-data/source/macron.png filter=lfs diff=lfs merge=lfs -text
-data/source/portrait1.png filter=lfs diff=lfs merge=lfs -text
-data/source/portrait2.png filter=lfs diff=lfs merge=lfs -text
-data/source/portrait3.png filter=lfs diff=lfs merge=lfs -text
-torch_extension/fused/fused_bias_act_kernel.cuda.o filter=lfs diff=lfs merge=lfs -text
-torch_extension/fused/fused_bias_act.o filter=lfs diff=lfs merge=lfs -text
-torch_extension/fused/fused.so filter=lfs diff=lfs merge=lfs -text
-torch_extension/upfirdn2d/upfirdn2d_kernel.cuda.o filter=lfs diff=lfs merge=lfs -text
-torch_extension/upfirdn2d/upfirdn2d.o filter=lfs diff=lfs merge=lfs -text
-torch_extension/upfirdn2d/upfirdn2d.so filter=lfs diff=lfs merge=lfs -text
-data/source/taylor.png filter=lfs diff=lfs merge=lfs -text
-data/driving/driving8.mp4 filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -4,9 +4,8 @@ emoji: 🐠
 colorFrom: pink
 colorTo: yellow
 sdk: gradio
-sdk_version: 5.41.1
+sdk_version: 5.42.0
 app_file: app.py
-python_version: 3.11
 pinned: false
 license: mit
 ---

app.py

@@ -1,21 +1,15 @@
 import gradio as gr
-import subprocess
-import os
-import spaces
 import torch
+from huggingface_hub import hf_hub_download
 
-extensions_dir = "./torch_extension/"
-os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
-
+from gradio_tabs.animation import animation
+from gradio_tabs.vid_edit import vid_edit
 from networks.generator import Generator
 
 device = torch.device("cuda")
 gen = Generator(size=512, motion_dim=40, scale=2).to(device)
-tmp_ckpt_path = "/home/user/.cache/torch/hub/checkpoints/lia-x.pt"
-if os.path.exists(tmp_ckpt_path):
-	gen.load_state_dict(torch.load(tmp_ckpt_path, weights_only=True))
-else:
-	gen.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/YaohuiW/LIA-X/resolve/main/lia-x.pt"))
+ckpt_path = hf_hub_download(repo_id="YaohuiW/LIA-X", filename="lia-x.pt")
+gen.load_state_dict(torch.load(ckpt_path, weights_only=True))
 gen.eval()
 
 chunk_size=8
@@ -38,18 +32,7 @@ custom_css = """
 </style>
 """
 
-
-# def load_tabs():
-# 	from gradio_tabs.animation import animation
-# 	from gradio_tabs.vid_edit import vid_edit
-    
-# 	animation()
-# 	vid_edit()
-
-
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-	# ... (input/output setup remains unchanged)
-
 	gr.HTML(load_file("assets/title.md"))
 	with gr.Row():
 		with gr.Accordion(open=False, label="Instruction"):
@@ -57,14 +40,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 	
 	with gr.Row():
 		with gr.Tabs():
-			from gradio_tabs.animation import animation
-			from gradio_tabs.vid_edit import vid_edit
 			animation(gen, chunk_size, device)
 			vid_edit(gen, chunk_size, device)
 
-    
-demo.launch(
-    server_name='0.0.0.0',
-    share=True,
-    allowed_paths=["./data/source","./data/driving"]
-)
+demo.launch(allowed_paths=["./data/source","./data/driving"])

gradio_tabs/animation.py

@@ -1,14 +1,12 @@
+import tempfile
+
 import gradio as gr
-import os
+import imageio
+import spaces
 import torch
 import torchvision
-from PIL import Image
 import numpy as np
-import imageio
-import spaces
-
-output_dir = "./res_gradio"
-os.makedirs(output_dir, exist_ok=True)
+from PIL import Image
 
 # lables
 labels_k = [
@@ -93,28 +91,32 @@ def vid_denorm(vid):
 	return vid
 
 
-def img_postprocessing(image, output_path=output_dir + "/output_img.png"):
+def img_postprocessing(image):
 
 	image = image.permute(0, 2, 3, 1)
 	edited_image = img_denorm(image)
 	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
-	imageio.imwrite(output_path, img_output, quality=6)
 
-	return output_path
+	with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+		imageio.imwrite(temp_file.name, img_output, quality=6)
+		return temp_file.name
+
 
 
-def vid_postprocessing(video, fps, output_path=output_dir + "/output_vid.mp4"):
+def vid_postprocessing(video, fps):
 	# video: BCTHW
 
 	vid = video.permute(0, 2, 3, 4, 1)	# B T H W C
 	vid_np = (vid_denorm(vid[0]).numpy() * 255).astype('uint8')
-	imageio.mimwrite(output_path, vid_np, fps=fps, codec='libx264', quality=10)
 
-	return output_path
+	with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
+		imageio.mimwrite(temp_file.name, vid_np, fps=fps, codec='libx264', quality=10)
+		return temp_file.name
 
 
 def animation(gen, chunk_size, device):
     
+	@spaces.GPU
 	@torch.no_grad()
 	def edit_media(image, *selected_s):
 
@@ -128,6 +130,7 @@ def animation(gen, chunk_size, device):
 
 		return edited_image
 
+	@spaces.GPU
 	@torch.no_grad()
 	def animate_media(image, video, *selected_s):
 
@@ -168,7 +171,6 @@ def animation(gen, chunk_size, device):
 								["./data/source/portrait3.png"],
 							],
 							inputs=[image_input],
-							cache_examples=False,
 							visible=True,
 							)
 
@@ -183,7 +185,6 @@ def animation(gen, chunk_size, device):
 								["./data/driving/driving8.mp4"],
 							],
 							inputs=[video_input],
-							cache_examples=False,
 							visible=True,
 							)
 
@@ -276,7 +277,6 @@ def animation(gen, chunk_size, device):
             fn=animate_media,
 			inputs=[image_input, video_input] + inputs_s,
             outputs=[image_output, video_output],
-            cache_examples=True,
 		)
 
 

gradio_tabs/vid_edit.py

@@ -1,15 +1,13 @@
+import tempfile
+
 import gradio as gr
-import os
+import imageio
+import numpy as np
+import spaces
 import torch
 import torchvision
-from PIL import Image
-import numpy as np
-import imageio
 from einops import rearrange
-import spaces
-
-output_dir = "./res_gradio"
-os.makedirs(output_dir, exist_ok=True)
+from PIL import Image
 
 # lables
 labels_k = [
@@ -94,16 +92,17 @@ def vid_denorm(vid):
 	return vid
 
 
-def img_postprocessing(image, output_path=output_dir + "/output_img.png"):
+def img_postprocessing(image):
 	image = image.permute(0, 2, 3, 1)
 	edited_image = img_denorm(image)
 	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
-	imageio.imwrite(output_path, img_output, quality=6)
 
-	return output_path
+	with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+		imageio.imwrite(temp_file.name, img_output, quality=6)
+		return temp_file.name
 
 
-def vid_all_save(vid_d, vid_a, fps, output_path=output_dir + "/output_vid.mp4", output_all_path=output_dir + "/output_all_vid.mp4"):
+def vid_all_save(vid_d, vid_a, fps):
 
 	vid_d = rearrange(vid_d, 'b t c h w -> b t h w c')
 	vid_a = rearrange(vid_a, 'b c t h w -> b t h w c')
@@ -112,15 +111,18 @@ def vid_all_save(vid_d, vid_a, fps, output_path=output_dir + "/output_vid.mp4",
 	vid_a_np = (vid_denorm(vid_a[0]).numpy() * 255).astype('uint8')
 	vid_all_np = (vid_denorm(vid_all[0]).numpy() * 255).astype('uint8')
 
-	imageio.mimwrite(output_path, vid_a_np, fps=fps, codec='libx264', quality=8)
-	imageio.mimwrite(output_all_path, vid_all_np, fps=fps, codec='libx264', quality=8)
+	with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as output_path:
+		imageio.mimwrite(output_path.name, vid_a_np, fps=fps, codec='libx264', quality=8)
 
-	return output_path, output_all_path
+	with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as output_all_path:
+		imageio.mimwrite(output_all_path.name, vid_all_np, fps=fps, codec='libx264', quality=8)
+
+	return output_path.name, output_all_path.name
 
 
 def vid_edit(gen, chunk_size, device):
     
-    
+	@spaces.GPU
 	@torch.no_grad()
 	def edit_img(video, *selected_s):
 
@@ -135,7 +137,7 @@ def vid_edit(gen, chunk_size, device):
 
 		return edited_image
 
-
+	@spaces.GPU
 	@torch.no_grad()
 	def edit_vid(video, *selected_s):
 
@@ -178,7 +180,6 @@ def vid_edit(gen, chunk_size, device):
 								["./data/driving/driving9.mp4"],
 							],
 							inputs=[video_input],
-							cache_examples=False,
 							visible=True,
 						)
 
@@ -272,7 +273,6 @@ def vid_edit(gen, chunk_size, device):
             fn=edit_vid,
 			inputs=[video_input] + inputs_s,
             outputs=[image_output, video_output, video_all_output],
-            cache_examples=True,
 		)
 
 

{torch_extension/fused → networks/op}/fused.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0434000b297e05aab7b7a7d537d32c3710cde943861c4cad5ef85db88964757d
-size 1297184
+oid sha256:a4d30ec1a330c9da4dd52648c85f278308a03d95e0cca21b9cbf29c538eb2e8b
+size 1433496

networks/op/fused_act.py

@@ -7,14 +7,17 @@ from torch.autograd import Function
 from torch.utils.cpp_extension import load
 
 
-module_path = os.path.dirname(__file__)
-fused = load(
-    "fused",
-    sources=[
-        os.path.join(module_path, "fused_bias_act.cpp"),
-        os.path.join(module_path, "fused_bias_act_kernel.cu"),
-    ],
-)
+if os.getenv("SPACE_ID"):
+    import networks.op.fused as fused
+else:
+    module_path = os.path.dirname(__file__)
+    fused = load(
+        "fused",
+        sources=[
+            os.path.join(module_path, "fused_bias_act.cpp"),
+            os.path.join(module_path, "fused_bias_act_kernel.cu"),
+        ],
+    )
 
 
 class FusedLeakyReLUFunctionBackward(Function):

networks/op/upfirdn2d.py

@@ -7,14 +7,17 @@ from torch.autograd import Function
 from torch.utils.cpp_extension import load
 
 
-module_path = os.path.dirname(__file__)
-upfirdn2d_op = load(
-    "upfirdn2d",
-    sources=[
-        os.path.join(module_path, "upfirdn2d.cpp"),
-        os.path.join(module_path, "upfirdn2d_kernel.cu"),
-    ],
-)
+if os.getenv("SPACE_ID"):
+    import networks.op.upfirdn2d_op as upfirdn2d_op
+else:
+    module_path = os.path.dirname(__file__)
+    upfirdn2d_op = load(
+        "upfirdn2d",
+        sources=[
+            os.path.join(module_path, "upfirdn2d.cpp"),
+            os.path.join(module_path, "upfirdn2d_kernel.cu"),
+        ],
+    )
 
 
 class UpFirDn2dBackward(Function):

torch_extension/upfirdn2d/upfirdn2d.o → networks/op/upfirdn2d_op.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9246c5a518be0a3b0909204f8b3a07bbccc9bf2ac6c413b15827f900d3452877
-size 1934968
+oid sha256:76d62a2db1f3493ccd4cd7f55c491d8d942926d45d68fae5a89af2a200029628
+size 1625920

requirements.txt

@@ -1,11 +1,7 @@
-torch==2.0.1
-torchvision==0.15.2
+av==15.0.0
 einops==0.8.1
 imageio==2.37.0
-numpy==1.24.3
-Pillow==11.3.0
-PyYAML==6.0.2
-tqdm==4.67.1
-ninja==1.11.1.4
-av==14.2.0
 imageio-ffmpeg==0.6.0
+ninja==1.11.1.4
+torch==2.5.1
+torchvision==0.20.1

torch_extension/fused/build.ninja

@@ -1,35 +0,0 @@
-ninja_required_version = 1.3
-cxx = c++
-nvcc = /mnt/petrelfs/share/cuda-11.7/bin/nvcc
-
-cflags = -DTORCH_EXTENSION_NAME=fused -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/TH -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/THC -isystem /mnt/petrelfs/share/cuda-11.7/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17
-post_cflags = 
-cuda_cflags = -DTORCH_EXTENSION_NAME=fused -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/TH -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/THC -isystem /mnt/petrelfs/share/cuda-11.7/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -std=c++17
-cuda_post_cflags = 
-cuda_dlink_post_cflags = 
-ldflags = -shared -L/mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/mnt/petrelfs/share/cuda-11.7/lib64 -lcudart
-
-rule compile
-  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
-  depfile = $out.d
-  deps = gcc
-
-rule cuda_compile
-  depfile = $out.d
-  deps = gcc
-  command = $nvcc  $cuda_cflags -c $in -o $out $cuda_post_cflags
-
-
-
-rule link
-  command = $cxx $in $ldflags -o $out
-
-build fused_bias_act.o: compile /mnt/hwfile/wangyaohui/projects/LIA-X/LIAX-release/networks/op/fused_bias_act.cpp
-build fused_bias_act_kernel.cuda.o: cuda_compile /mnt/hwfile/wangyaohui/projects/LIA-X/LIAX-release/networks/op/fused_bias_act_kernel.cu
-
-
-
-build fused.so: link fused_bias_act.o fused_bias_act_kernel.cuda.o
-
-default fused.so
-

torch_extension/fused/fused_bias_act.o

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:11ef1d0528905c05b8866df6448a36da47e96bdf2c52e840ddd5c876791dc1e6
-size 1928096

torch_extension/fused/fused_bias_act_kernel.cuda.o

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4a576cabbb56d980122af7c924410f8160fa3c81be957f1a05b2c75a932252a5
-size 579704

torch_extension/upfirdn2d/build.ninja

@@ -1,35 +0,0 @@
-ninja_required_version = 1.3
-cxx = c++
-nvcc = /mnt/petrelfs/share/cuda-11.7/bin/nvcc
-
-cflags = -DTORCH_EXTENSION_NAME=upfirdn2d -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/TH -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/THC -isystem /mnt/petrelfs/share/cuda-11.7/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17
-post_cflags = 
-cuda_cflags = -DTORCH_EXTENSION_NAME=upfirdn2d -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/TH -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/include/THC -isystem /mnt/petrelfs/share/cuda-11.7/include -isystem /mnt/petrelfs/wangyaohui/anaconda3/envs/latte/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -std=c++17
-cuda_post_cflags = 
-cuda_dlink_post_cflags = 
-ldflags = -shared -L/mnt/petrelfs/wangyaohui/anaconda3/envs/latte/lib/python3.11/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/mnt/petrelfs/share/cuda-11.7/lib64 -lcudart
-
-rule compile
-  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
-  depfile = $out.d
-  deps = gcc
-
-rule cuda_compile
-  depfile = $out.d
-  deps = gcc
-  command = $nvcc  $cuda_cflags -c $in -o $out $cuda_post_cflags
-
-
-
-rule link
-  command = $cxx $in $ldflags -o $out
-
-build upfirdn2d.o: compile /mnt/hwfile/wangyaohui/projects/LIA-X/LIAX-release/networks/op/upfirdn2d.cpp
-build upfirdn2d_kernel.cuda.o: cuda_compile /mnt/hwfile/wangyaohui/projects/LIA-X/LIAX-release/networks/op/upfirdn2d_kernel.cu
-
-
-
-build upfirdn2d.so: link upfirdn2d.o upfirdn2d_kernel.cuda.o
-
-default upfirdn2d.so
-

torch_extension/upfirdn2d/upfirdn2d.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:277ef4534d19f242b1ee96cdbaf08b638f3fa5f2beadf094c0399c0e17d3990a
-size 1390992

torch_extension/upfirdn2d/upfirdn2d_kernel.cuda.o

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1d36964bd9f9c07018f4b4761a637b3846c2b5ea6247483ed243b8ed1cff3a2f
-size 658616