# Copyright (c) 2024-2025, Yisheng He, Yuan Dong
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
os.system("rm -rf /data-nvme/zerogpu-offload/")
os.system("pip install chumpy")
# os.system("pip uninstall -y basicsr")
os.system("pip install Cython")
os.system("pip install ./new_wheels/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl")
os.system("pip install ./wheels/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl")
os.system("pip install ./wheels/nvdiffrast-0.3.3-cp310-cp310-linux_x86_64.whl --force-reinstall")
os.system(
"pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt240/download.html")
os.system("pip install numpy==1.23.0")
import cv2
import sys
import base64
import subprocess
import argparse
from glob import glob
import gradio as gr
import numpy as np
from PIL import Image
from omegaconf import OmegaConf
import torch
import moviepy.editor as mpy
from lam.runners.infer.head_utils import prepare_motion_seqs, preprocess_image
from lam.utils.ffmpeg_utils import images_to_video
import spaces
def compile_module(subfolder, script):
try:
# Save the current working directory
current_dir = os.getcwd()
# Change directory to the subfolder
os.chdir(os.path.join(current_dir, subfolder))
# Run the compilation command
result = subprocess.run(
["sh", script],
capture_output=True,
text=True,
check=True
)
# Print the compilation output
print("Compilation output:", result.stdout)
except Exception as e:
# Print any error that occurred
print(f"An error occurred: {e}")
finally:
# Ensure returning to the original directory
os.chdir(current_dir)
print("Returned to the original directory.")
# compile flame_tracking dependence submodule
compile_module("external/landmark_detection/FaceBoxesV2/utils/", "make.sh")
from flame_tracking_single_image import FlameTrackingSingleImage
def launch_pretrained():
from huggingface_hub import snapshot_download, hf_hub_download
# launch pretrained for flame tracking.
hf_hub_download(repo_id='yuandong513/flametracking_model',
repo_type='model',
filename='pretrain_model.tar',
local_dir='./')
os.system('tar -xf pretrain_model.tar && rm pretrain_model.tar')
# launch human model files
hf_hub_download(repo_id='3DAIGC/LAM-assets',
repo_type='model',
filename='LAM_human_model.tar',
local_dir='./')
os.system('tar -xf LAM_human_model.tar && rm LAM_human_model.tar')
# launch pretrained for LAM
model_dir = hf_hub_download(repo_id="3DAIGC/LAM-20K", repo_type="model", local_dir="./exps/releases/lam/lam-20k/step_045500/", filename="config.json")
print(model_dir)
model_dir = hf_hub_download(repo_id="3DAIGC/LAM-20K", repo_type="model", local_dir="./exps/releases/lam/lam-20k/step_045500/", filename="model.safetensors")
print(model_dir)
model_dir = hf_hub_download(repo_id="3DAIGC/LAM-20K", repo_type="model", local_dir="./exps/releases/lam/lam-20k/step_045500/", filename="README.md")
print(model_dir)
# launch example for LAM
hf_hub_download(repo_id='3DAIGC/LAM-assets',
repo_type='model',
filename='LAM_assets.tar',
local_dir='./')
os.system('tar -xf LAM_assets.tar && rm LAM_assets.tar')
hf_hub_download(repo_id='3DAIGC/LAM-assets',
repo_type='model',
filename='config.json',
local_dir='./tmp/')
def launch_env_not_compile_with_cuda():
os.system('pip install chumpy')
os.system('pip install numpy==1.23.0')
os.system(
'pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt251/download.html'
)
def assert_input_image(input_image):
if input_image is None:
raise gr.Error('No image selected or uploaded!')
def prepare_working_dir():
import tempfile
working_dir = tempfile.TemporaryDirectory()
return working_dir
def init_preprocessor():
from lam.utils.preprocess import Preprocessor
global preprocessor
preprocessor = Preprocessor()
def preprocess_fn(image_in: np.ndarray, remove_bg: bool, recenter: bool,
working_dir):
image_raw = os.path.join(working_dir.name, 'raw.png')
with Image.fromarray(image_in) as img:
img.save(image_raw)
image_out = os.path.join(working_dir.name, 'rembg.png')
success = preprocessor.preprocess(image_path=image_raw,
save_path=image_out,
rmbg=remove_bg,
recenter=recenter)
assert success, f'Failed under preprocess_fn!'
return image_out
def get_image_base64(path):
with open(path, 'rb') as image_file:
encoded_string = base64.b64encode(image_file.read()).decode()
return f'data:image/png;base64,{encoded_string}'
def save_imgs_2_video(imgs, v_pth, fps=30):
# moviepy example
from moviepy.editor import ImageSequenceClip, VideoFileClip
images = [image.astype(np.uint8) for image in imgs]
clip = ImageSequenceClip(images, fps=fps)
# final_duration = len(images) / fps
# clip = clip.subclip(0, final_duration)
clip = clip.subclip(0, len(images) / fps)
clip.write_videofile(v_pth, codec='libx264')
import cv2
cap = cv2.VideoCapture(v_pth)
nf = cap.get(cv2.CAP_PROP_FRAME_COUNT)
if nf != len(images):
print("="*100+f"\n{v_pth} moviepy saved video frame error."+"\n"+"="*100)
print(f"Video saved successfully at {v_pth}")
def add_audio_to_video(video_path, out_path, audio_path, fps=30):
# Import necessary modules from moviepy
from moviepy.editor import VideoFileClip, AudioFileClip
# Load video file into VideoFileClip object
video_clip = VideoFileClip(video_path)
# Load audio file into AudioFileClip object
audio_clip = AudioFileClip(audio_path)
# Hard code clip audio
if audio_clip.duration > 10:
audio_clip = audio_clip.subclip(0, 10)
# Attach audio clip to video clip (replaces existing audio)
video_clip_with_audio = video_clip.set_audio(audio_clip)
# Export final video with audio using standard codecs
video_clip_with_audio.write_videofile(out_path, codec='libx264', audio_codec='aac', fps=fps)
print(f"Audio added successfully at {out_path}")
def parse_configs():
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str)
parser.add_argument("--infer", type=str)
args, unknown = parser.parse_known_args()
cfg = OmegaConf.create()
cli_cfg = OmegaConf.from_cli(unknown)
# parse from ENV
if os.environ.get("APP_INFER") is not None:
args.infer = os.environ.get("APP_INFER")
if os.environ.get("APP_MODEL_NAME") is not None:
cli_cfg.model_name = os.environ.get("APP_MODEL_NAME")
args.config = args.infer if args.config is None else args.config
if args.config is not None:
cfg_train = OmegaConf.load(args.config)
cfg.source_size = cfg_train.dataset.source_image_res
try:
cfg.src_head_size = cfg_train.dataset.src_head_size
except:
cfg.src_head_size = 112
cfg.render_size = cfg_train.dataset.render_image.high
_relative_path = os.path.join(
cfg_train.experiment.parent,
cfg_train.experiment.child,
os.path.basename(cli_cfg.model_name).split("_")[-1],
)
cfg.save_tmp_dump = os.path.join("exps", "save_tmp", _relative_path)
cfg.image_dump = os.path.join("exps", "images", _relative_path)
cfg.video_dump = os.path.join("exps", "videos", _relative_path) # output path
if args.infer is not None:
cfg_infer = OmegaConf.load(args.infer)
cfg.merge_with(cfg_infer)
cfg.setdefault(
"save_tmp_dump", os.path.join("exps", cli_cfg.model_name, "save_tmp")
)
cfg.setdefault("image_dump", os.path.join("exps", cli_cfg.model_name, "images"))
cfg.setdefault(
"video_dump", os.path.join("dumps", cli_cfg.model_name, "videos")
)
cfg.setdefault("mesh_dump", os.path.join("dumps", cli_cfg.model_name, "meshes"))
cfg.motion_video_read_fps = 30
cfg.merge_with(cli_cfg)
cfg.setdefault("logger", "INFO")
assert cfg.model_name is not None, "model_name is required"
return cfg, cfg_train
def demo_lam(flametracking, lam, cfg):
@spaces.GPU(duration=80)
def core_fn(image_path: str, video_params, working_dir):
image_raw = os.path.join(working_dir.name, "raw.png")
with Image.open(image_path).convert('RGB') as img:
img.save(image_raw)
base_vid = os.path.basename(video_params).split(".")[0]
flame_params_dir = os.path.join("./assets/sample_motion/export", base_vid, "flame_param")
base_iid = os.path.basename(image_path).split('.')[0]
image_path = os.path.join("./assets/sample_input", base_iid, "images/00000_00.png")
dump_video_path = os.path.join(working_dir.name, "output.mp4")
dump_image_path = os.path.join(working_dir.name, "output.png")
# prepare dump paths
omit_prefix = os.path.dirname(image_raw)
image_name = os.path.basename(image_raw)
uid = image_name.split(".")[0]
subdir_path = os.path.dirname(image_raw).replace(omit_prefix, "")
subdir_path = (
subdir_path[1:] if subdir_path.startswith("/") else subdir_path
)
print("subdir_path and uid:", subdir_path, uid)
motion_seqs_dir = flame_params_dir
dump_image_dir = os.path.dirname(dump_image_path)
os.makedirs(dump_image_dir, exist_ok=True)
print(image_raw, motion_seqs_dir, dump_image_dir, dump_video_path)
dump_tmp_dir = dump_image_dir
if os.path.exists(dump_video_path):
return dump_image_path, dump_video_path
motion_img_need_mask = cfg.get("motion_img_need_mask", False) # False
vis_motion = cfg.get("vis_motion", False) # False
# preprocess input image: segmentation, flame params estimation
# """
return_code = flametracking.preprocess(image_raw)
assert (return_code == 0), "flametracking preprocess failed!"
return_code = flametracking.optimize()
assert (return_code == 0), "flametracking optimize failed!"
return_code, output_dir = flametracking.export()
assert (return_code == 0), "flametracking export failed!"
image_path = os.path.join(output_dir, "images/00000_00.png")
# """
mask_path = image_path.replace("/images/", "/fg_masks/").replace(".jpg", ".png")
print(image_path, mask_path)
aspect_standard = 1.0 / 1.0
source_size = cfg.source_size
render_size = cfg.render_size
render_fps = 30
# prepare reference image
image, _, _, shape_param = preprocess_image(image_path, mask_path=mask_path, intr=None, pad_ratio=0,
bg_color=1.,
max_tgt_size=None, aspect_standard=aspect_standard,
enlarge_ratio=[1.0, 1.0],
render_tgt_size=source_size, multiply=14, need_mask=True,
get_shape_param=True)
# save masked image for vis
save_ref_img_path = os.path.join(dump_tmp_dir, "output.png")
vis_ref_img = (image[0].permute(1, 2, 0).cpu().detach().numpy() * 255).astype(np.uint8)
Image.fromarray(vis_ref_img).save(save_ref_img_path)
# prepare motion seq
src = image_path.split('/')[-3]
driven = motion_seqs_dir.split('/')[-2]
src_driven = [src, driven]
motion_seq = prepare_motion_seqs(motion_seqs_dir, None, save_root=dump_tmp_dir, fps=render_fps,
bg_color=1., aspect_standard=aspect_standard, enlarge_ratio=[1.0, 1, 0],
render_image_res=render_size, multiply=16,
need_mask=motion_img_need_mask, vis_motion=vis_motion,
shape_param=shape_param, test_sample=False, cross_id=False,
src_driven=src_driven, max_squen_length=300)
# start inference
motion_seq["flame_params"]["betas"] = shape_param.unsqueeze(0)
device, dtype = "cuda", torch.float32
print("start to inference...................")
with torch.no_grad():
# TODO check device and dtype
res = lam.infer_single_view(image.unsqueeze(0).to(device, dtype), None, None,
render_c2ws=motion_seq["render_c2ws"].to(device),
render_intrs=motion_seq["render_intrs"].to(device),
render_bg_colors=motion_seq["render_bg_colors"].to(device),
flame_params={k: v.to(device) for k, v in motion_seq["flame_params"].items()})
rgb = res["comp_rgb"].detach().cpu().numpy() # [Nv, H, W, 3], 0-1
mask = res["comp_mask"].detach().cpu().numpy() # [Nv, H, W, 3], 0-1
mask[mask < 0.5] = 0.0
rgb = rgb * mask + (1 - mask) * 1
rgb = (np.clip(rgb, 0, 1.0) * 255).astype(np.uint8)
if vis_motion:
vis_ref_img = np.tile(
cv2.resize(vis_ref_img, (rgb[0].shape[1], rgb[0].shape[0]), interpolation=cv2.INTER_AREA)[None, :, :,
:],
(rgb.shape[0], 1, 1, 1),
)
rgb = np.concatenate([vis_ref_img, rgb, motion_seq["vis_motion_render"]], axis=2)
os.makedirs(os.path.dirname(dump_video_path), exist_ok=True)
print("==="*36, "\nrgb length:", rgb.shape, render_fps, "==="*36)
save_imgs_2_video(rgb, dump_video_path, render_fps)
# images_to_video(rgb, output_path=dump_video_path, fps=30, gradio_codec=False, verbose=True)
audio_path = os.path.join("./assets/sample_motion/export", base_vid, base_vid + ".wav")
dump_video_path_wa = dump_video_path.replace(".mp4", "_audio.mp4")
add_audio_to_video(dump_video_path, dump_video_path_wa, audio_path)
return dump_image_path, dump_video_path_wa
def core_fn_space(image_path: str, video_params, working_dir):
return core_fn(image_path, video_params, working_dir)
with gr.Blocks(analytics_enabled=False, delete_cache=[3600, 3600]) as demo:
logo_url = './assets/images/logo.jpeg'
logo_base64 = get_image_base64(logo_url)
gr.HTML(f"""
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
<div>
<h1> <img src="{logo_base64}" style='height:35px; display:inline-block;'/> Large Avatar Model for One-shot Animatable Gaussian Head</h1>
</div>
</div>
""")
gr.HTML(
"""
<div style="display: flex; justify-content: center; align-items: center; text-align: center; margin: 20px; gap: 10px;">
<a class="flex-item" href="https://arxiv.org/abs/2502.17796" target="_blank">
<img src="https://img.shields.io/badge/Paper-arXiv-darkred.svg" alt="arXiv Paper">
</a>
<a class="flex-item" href="https://aigc3d.github.io/projects/LAM/" target="_blank">
<img src="https://img.shields.io/badge/Project-LAM-blue" alt="Project Page">
</a>
<a class="flex-item" href="https://github.com/aigc3d/LAM" target="_blank">
<img src="https://img.shields.io/github/stars/aigc3d/LAM?label=Github%20★&logo=github&color=C8C" alt="badge-github-stars">
</a>
<a class="flex-item" href="https://youtu.be/FrfE3RYSKhk" target="_blank">
<img src="https://img.shields.io/badge/Youtube-Video-red.svg" alt="Video">
</a>
</div>
"""
)
gr.HTML("""<div style="margin-top: -10px">
<p style="margin: 4px 0; line-height: 1.2"><h4 style="color: black; margin: 2px 0">Notes1: Inputing front-face images or face orientation close to the driven signal gets better results.</h4></p>
<p style="margin: 4px 0; line-height: 1.2"><h4 style="color: black; margin: 2px 0">Notes2: Due to computational constraints with Hugging Face's ZeroGPU infrastructure, 3D avatar generation requires ~1 minute per instance.</h4></p>
<p style="margin: 4px 0; line-height: 1.2"><h4 style="color: black; margin: 2px 0">Notes3: Using LAM-20K model (lower quality than premium LAM-80K) to mitigate processing latency.</h4></p>
</div>""")
# DISPLAY
with gr.Row():
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id='lam_input_image'):
with gr.TabItem('Input Image'):
with gr.Row():
input_image = gr.Image(label='Input Image',
image_mode='RGB',
height=480,
width=270,
sources='upload',
type='filepath',
elem_id='content_image')
# EXAMPLES
with gr.Row():
examples = [
['assets/sample_input/messi.png'],
['assets/sample_input/status.png'],
['assets/sample_input/james.png'],
['assets/sample_input/cluo.jpg'],
['assets/sample_input/dufu.jpg'],
['assets/sample_input/libai.jpg'],
['assets/sample_input/barbara.jpg'],
['assets/sample_input/pop.png'],
['assets/sample_input/musk.jpg'],
['assets/sample_input/speed.jpg'],
['assets/sample_input/zhouxingchi.jpg'],
]
gr.Examples(
examples=examples,
inputs=[input_image],
examples_per_page=20
)
with gr.Column():
with gr.Tabs(elem_id='lam_input_video'):
with gr.TabItem('Input Video'):
with gr.Row():
video_input = gr.Video(label='Input Video',
height=480,
width=270,
interactive=False)
examples = ['./assets/sample_motion/export/Speeding_Scandal/Speeding_Scandal.mp4',
'./assets/sample_motion/export/Look_In_My_Eyes/Look_In_My_Eyes.mp4',
'./assets/sample_motion/export/D_ANgelo_Dinero/D_ANgelo_Dinero.mp4',
'./assets/sample_motion/export/Michael_Wayne_Rosen/Michael_Wayne_Rosen.mp4',
'./assets/sample_motion/export/I_Am_Iron_Man/I_Am_Iron_Man.mp4',
'./assets/sample_motion/export/Anti_Drugs/Anti_Drugs.mp4',
'./assets/sample_motion/export/Pen_Pineapple_Apple_Pen/Pen_Pineapple_Apple_Pen.mp4',
'./assets/sample_motion/export/Joe_Biden/Joe_Biden.mp4',
'./assets/sample_motion/export/Donald_Trump/Donald_Trump.mp4',
'./assets/sample_motion/export/Taylor_Swift/Taylor_Swift.mp4',
'./assets/sample_motion/export/GEM/GEM.mp4',
'./assets/sample_motion/export/The_Shawshank_Redemption/The_Shawshank_Redemption.mp4'
]
print("Video example list {}".format(examples))
gr.Examples(
examples=examples,
inputs=[video_input],
examples_per_page=20,
)
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id='lam_processed_image'):
with gr.TabItem('Processed Image'):
with gr.Row():
processed_image = gr.Image(
label='Processed Image',
image_mode='RGBA',
type='filepath',
elem_id='processed_image',
height=480,
width=270,
interactive=False)
with gr.Column(variant='panel', scale=1):
with gr.Tabs(elem_id='lam_render_video'):
with gr.TabItem('Rendered Video'):
with gr.Row():
output_video = gr.Video(label='Rendered Video',
format='mp4',
height=480,
width=270,
autoplay=True)
# SETTING
with gr.Row():
with gr.Column(variant='panel', scale=1):
submit = gr.Button('Generate',
elem_id='lam_generate',
variant='primary')
main_fn = core_fn
working_dir = gr.State()
submit.click(
fn=assert_input_image,
inputs=[input_image],
queue=False,
).success(
fn=prepare_working_dir,
outputs=[working_dir],
queue=False,
).success(
fn=main_fn,
inputs=[input_image, video_input,
working_dir], # video_params refer to smpl dir
outputs=[processed_image, output_video],
)
demo.queue()
demo.launch()
def _build_model(cfg):
from lam.models import model_dict
from lam.utils.hf_hub import wrap_model_hub
hf_model_cls = wrap_model_hub(model_dict["lam"])
model = hf_model_cls.from_pretrained(cfg.model_name)
return model
def launch_gradio_app():
os.environ.update({
'APP_ENABLED': '1',
'APP_MODEL_NAME':
'./exps/releases/lam/lam-20k/step_045500/',
'APP_INFER': './configs/inference/lam-20k-8gpu.yaml',
'APP_TYPE': 'infer.lam',
'NUMBA_THREADING_LAYER': 'omp',
})
cfg, _ = parse_configs()
lam = _build_model(cfg)
lam.to('cuda')
flametracking = FlameTrackingSingleImage(output_dir='tracking_output',
alignment_model_path='./pretrain_model/68_keypoints_model.pkl',
vgghead_model_path='./pretrain_model/vgghead/vgg_heads_l.trcd',
human_matting_path='./pretrain_model/matting/stylematte_synth.pt',
facebox_model_path='./pretrain_model/FaceBoxesV2.pth',
detect_iris_landmarks=False)
demo_lam(flametracking, lam, cfg)
if __name__ == '__main__':
launch_pretrained()
launch_gradio_app()