Create app.py

app.py

@@ -0,0 +1,317 @@
+import gradio as gr
+import torch
+import os
+import numpy as np
+from PIL import Image
+import cv2
+import tempfile
+import moviepy.editor as mp
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+from diffusers.utils import export_to_video, load_image
+
+# Import required modules from SkyReels
+from skyreels_a1.models.transformer3d import CogVideoXTransformer3DModel
+from skyreels_a1.skyreels_a1_i2v_pipeline import SkyReelsA1ImagePoseToVideoPipeline
+from skyreels_a1.pre_process_lmk3d import FaceAnimationProcessor
+from skyreels_a1.src.media_pipe.mp_utils import LMKExtractor
+from skyreels_a1.src.media_pipe.draw_util_2d import FaceMeshVisualizer2d
+
+from diffusers.models import AutoencoderKLCogVideoX
+from transformers import SiglipImageProcessor, SiglipVisionModel
+from diffposetalk.diffposetalk import DiffPoseTalk
+
+from huggingface_hub import snapshot_download
+
+os.system("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt221/download.html")
+
+os.makedirs("pretrained_models", exist_ok=True)
+
+snapshot_download(
+    repo_id="multimodalart/diffposetalk",
+    local_dir="pretrained_models/diffposetalk"
+)
+
+snapshot_download(
+    repo_id="Skywork/SkyReels-A1",
+    local_dir="pretrained_models/FLAME",
+    allow_patterns="extra_models/FLAME/**"
+)
+
+snapshot_download(
+    repo_id="Skywork/SkyReels-A1",
+    local_dir="pretrained_models/mediapipe",
+    allow_patterns="extra_models/mediapipe/**"
+)
+
+snapshot_download(
+    repo_id="Skywork/SkyReels-A1",
+    local_dir="pretrained_models/smirk",
+    allow_patterns="extra_models/smirk/**"
+)
+
+snapshot_download(
+    repo_id="Skywork/SkyReels-A1",
+    local_dir="pretrained_models/SkyReels-A1-5B",
+    allow_patterns="SkyReels-A1-5B/**"
+)
+
+
+# Helper functions from the original script
+def parse_video(driving_frames, max_frame_num, fps=25):
+    video_length = len(driving_frames)
+    duration = video_length / fps 
+    target_times = np.arange(0, duration, 1/12)
+    frame_indices = (target_times * fps).astype(np.int32)
+
+    frame_indices = frame_indices[frame_indices < video_length]
+    new_driving_frames = []
+    for idx in frame_indices:
+        new_driving_frames.append(driving_frames[idx])
+        if len(new_driving_frames) >= max_frame_num - 1:
+            break
+
+    video_lenght_add = max_frame_num - len(new_driving_frames) - 1
+    new_driving_frames = [new_driving_frames[0]]*2 + new_driving_frames[1:len(new_driving_frames)-1] + [new_driving_frames[-1]] * video_lenght_add
+    return new_driving_frames
+
+def write_mp4(video_path, samples, fps=12):
+    clip = mp.ImageSequenceClip(samples, fps=fps)
+    clip.write_videofile(video_path, audio_codec="aac", codec="libx264",
+                         ffmpeg_params=["-crf", "18", "-preset", "slow"])
+
+def save_video_with_audio(video_path, audio_path, save_path):
+    video_clip = mp.VideoFileClip(video_path)
+    audio_clip = mp.AudioFileClip(audio_path)
+
+    if audio_clip.duration > video_clip.duration:
+        audio_clip = audio_clip.subclip(0, video_clip.duration)
+    
+    video_with_audio = video_clip.set_audio(audio_clip)
+    video_with_audio.write_videofile(save_path, fps=12, codec="libx264", audio_codec="aac")
+    
+    # Clean up
+    video_clip.close()
+    audio_clip.close()
+    return save_path
+
+# Global parameters
+model_name = "pretrained_models/SkyReels-A1-5B/"
+siglip_name = "pretrained_models/SkyReels-A1-5B/siglip-so400m-patch14-384"
+weight_dtype = torch.bfloat16
+max_frame_num = 49
+sample_size = [480, 720]
+
+# Preload all models in global context
+print("Loading models...")
+
+# Load LMK extractor and processors
+lmk_extractor = LMKExtractor()
+processor = FaceAnimationProcessor(checkpoint='pretrained_models/smirk/SMIRK_em1.pt')
+vis = FaceMeshVisualizer2d(forehead_edge=False, draw_head=False, draw_iris=False)
+face_helper = FaceRestoreHelper(upscale_factor=1, face_size=512, crop_ratio=(1, 1), 
+                            det_model='retinaface_resnet50', save_ext='png', device="cuda")
+
+# Load siglip visual encoder
+siglip = SiglipVisionModel.from_pretrained(siglip_name)
+siglip_normalize = SiglipImageProcessor.from_pretrained(siglip_name)
+
+# Load diffposetalk
+diffposetalk = DiffPoseTalk()
+
+# Load SkyReels models
+transformer = CogVideoXTransformer3DModel.from_pretrained(
+    model_name, 
+    subfolder="transformer"
+).to(weight_dtype)
+
+vae = AutoencoderKLCogVideoX.from_pretrained(
+    model_name, 
+    subfolder="vae"
+).to(weight_dtype)
+
+lmk_encoder = AutoencoderKLCogVideoX.from_pretrained(
+    model_name, 
+    subfolder="pose_guider",
+).to(weight_dtype)
+
+# Set up pipeline
+pipe = SkyReelsA1ImagePoseToVideoPipeline.from_pretrained(
+    model_name,
+    transformer=transformer,
+    vae=vae,
+    lmk_encoder=lmk_encoder,
+    image_encoder=siglip, 
+    feature_extractor=siglip_normalize,
+    torch_dtype=torch.bfloat16
+)
+pipe.to("cuda")
+pipe.transformer = torch.compile(pipe.transformer) 
+pipe.vae = torch.compile(pipe.vae)
+# pipe.enable_model_cpu_offload()
+# pipe.vae.enable_tiling()
+
+print("Models loaded successfully!")
+
+def process_image_audio(image_path, audio_path, guidance_scale=3.0, steps=10, progress=gr.Progress()):
+    progress(0.1, desc="Processing inputs...")
+    # Create a directory for outputs if it doesn't exist
+    output_dir = "gradio_outputs"
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Create temp files for processing
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video_file, \
+         tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_output_file:
+        temp_video_path = temp_video_file.name
+        final_output_path = temp_output_file.name
+    
+    # Set seed
+    seed = 43
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    
+    progress(0.2, desc="Processing image...")
+    # Load and process image
+    image = load_image(image=image_path)
+    image = processor.crop_and_resize(image, sample_size[0], sample_size[1])
+    
+    # Crop face
+    ref_image, x1, y1 = processor.face_crop(np.array(image))
+    face_h, face_w, _ = ref_image.shape
+    source_image = ref_image
+    
+    progress(0.3, desc="Processing facial landmarks...")
+    # Process source image
+    source_outputs, source_tform, image_original = processor.process_source_image(source_image)
+    
+    progress(0.4, desc="Processing audio...")
+    # Process audio and generate driving outputs
+    driving_outputs = diffposetalk.infer_from_file(
+        audio_path, 
+        source_outputs["shape_params"].view(-1)[:100].detach().cpu().numpy()
+    )
+    
+    progress(0.5, desc="Processing landmarks from coefficients...")
+    # Process landmarks
+    out_frames = processor.preprocess_lmk3d_from_coef(
+        source_outputs, source_tform, image_original.shape, driving_outputs
+    )
+    out_frames = parse_video(out_frames, max_frame_num)
+    
+    rescale_motions = np.zeros_like(image)[np.newaxis, :].repeat(48, axis=0)
+    for ii in range(rescale_motions.shape[0]):
+        rescale_motions[ii][y1:y1+face_h, x1:x1+face_w] = out_frames[ii]
+    
+    ref_image_resized = cv2.resize(ref_image, (512, 512))
+    ref_lmk = lmk_extractor(ref_image_resized[:, :, ::-1])
+    
+    ref_img = vis.draw_landmarks_v3(
+        (512, 512), (face_w, face_h), 
+        ref_lmk['lmks'].astype(np.float32), normed=True
+    )
+    
+    first_motion = np.zeros_like(np.array(image))
+    first_motion[y1:y1+face_h, x1:x1+face_w] = ref_img
+    first_motion = first_motion[np.newaxis, :]
+    
+    motions = np.concatenate([first_motion, rescale_motions])
+    input_video = motions[:max_frame_num]
+    
+    # Face alignment
+    face_helper.clean_all() 
+    face_helper.read_image(np.array(image)[:, :, ::-1])
+    face_helper.get_face_landmarks_5(only_center_face=True)
+    face_helper.align_warp_face()
+    align_face = face_helper.cropped_faces[0]
+    image_face = align_face[:, :, ::-1]
+    
+    # Prepare input video
+    input_video = torch.from_numpy(np.array(input_video)).permute([3, 0, 1, 2]).unsqueeze(0)
+    input_video = input_video / 255
+    
+    progress(0.6, desc="Generating animation (this may take a while)...")
+    # Generate video
+    #with torch.no_grad():
+    sample = pipe(
+        image=image,
+        image_face=image_face,
+        control_video=input_video,
+        prompt="", 
+        negative_prompt="",
+        height=480,
+        width=720,
+        num_frames=49,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        num_inference_steps=steps,
+    )
+    out_samples = sample.frames[0]
+    
+    out_samples = out_samples[2:]  # Skip first two frames
+    
+    progress(0.8, desc="Creating output video...")
+    # Export video
+    export_to_video(out_samples, temp_video_path, fps=12)
+    
+    progress(0.9, desc="Adding audio to video...")
+    # Add audio to video
+    result_path = save_video_with_audio(temp_video_path, audio_path, final_output_path)
+    
+    # Create side-by-side comparison
+    target_h, target_w = sample_size[0], sample_size[1]
+    final_images = []
+    for i in range(len(out_samples)):
+        frame1 = image
+        frame2 = Image.fromarray(np.array(out_samples[i])).convert("RGB")
+        
+        result = Image.new('RGB', (target_w * 2, target_h))
+        result.paste(frame1, (0, 0))
+        result.paste(frame2, (target_w, 0))
+        final_images.append(np.array(result))
+    
+    comparison_path = os.path.join(output_dir, "comparison.mp4") 
+    write_mp4(comparison_path, final_images, fps=12)
+    
+    # Add audio to comparison video
+    comparison_with_audio = os.path.join(output_dir, "comparison_with_audio.mp4")
+    comparison_with_audio = save_video_with_audio(comparison_path, audio_path, comparison_with_audio)
+    
+    progress(1.0, desc="Done!")
+    return result_path, comparison_with_audio
+
+# Create Gradio interface
+with gr.Blocks(title="SkyReels A1 Face Animation") as app:
+    gr.Markdown("# SkyReels A1 Face Animation")
+    gr.Markdown("Upload a portrait image and an audio file to animate the face")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="filepath", label="Portrait Image")
+            audio_input = gr.Audio(type="filepath", label="Driving Audio")
+            
+            with gr.Row():
+                guidance_scale = gr.Slider(minimum=1.0, maximum=7.0, value=3.0, step=0.1, label="Guidance Scale")
+                inference_steps = gr.Slider(minimum=5, maximum=30, value=10, step=1, label="Inference Steps")
+            
+            generate_button = gr.Button("Generate Animation", variant="primary")
+        
+        with gr.Column():
+            output_video = gr.Video(label="Animation Result")
+            comparison_video = gr.Video(label="Side-by-Side Comparison")
+    
+    generate_button.click(
+        fn=process_image_audio,
+        inputs=[image_input, audio_input, guidance_scale, inference_steps],
+        outputs=[output_video, comparison_video]
+    )
+    
+    gr.Markdown("""
+    ## Instructions
+    1. Upload a portrait image (frontal face works best)
+    2. Upload an audio file (wav format recommended)
+    3. Adjust parameters if needed
+    4. Click "Generate Animation" to create the video
+    
+    Note: Processing may take several minutes depending on your hardware.
+    """)
+
+if __name__ == "__main__":
+    app.launch(share=True)