Spaces:

VIDraft
/

Portrait-Animation

Running on Zero

File size: 15,780 Bytes

import spaces
import gradio as gr
import os
import numpy as np
from pydub import AudioSegment
import hashlib
import io
from sonic import Sonic
from PIL import Image
import torch

# 초기 실행 시 필요한 모델들을 다운로드
cmd = (
    'python3 -m pip install "huggingface_hub[cli]"; '
    'huggingface-cli download LeonJoe13/Sonic --local-dir checkpoints; '
    'huggingface-cli download stabilityai/stable-video-diffusion-img2vid-xt --local-dir  checkpoints/stable-video-diffusion-img2vid-xt; '
    'huggingface-cli download openai/whisper-tiny --local-dir checkpoints/whisper-tiny;'
)
os.system(cmd)

pipe = Sonic()

def get_md5(content_bytes):
    """MD5 해시를 계산하여 32자리 문자열을 반환"""
    return hashlib.md5(content_bytes).hexdigest()

tmp_path = './tmp_path/'
res_path = './res_path/'
os.makedirs(tmp_path, exist_ok=True)
os.makedirs(res_path, exist_ok=True)


@spaces.GPU(duration=600)  # 긴 비디오 처리를 위해 duration 600초로 설정 (10분)
def get_video_res(img_path, audio_path, res_video_path, dynamic_scale=1.0):
    """
    Sonic pipeline으로부터 실제 비디오를 생성하는 함수.
    최대 60초 길이의 오디오에 대해 inference_steps를 결정하여,
    얼굴 탐지 후 영상 생성 작업을 수행함.
    """
    expand_ratio = 0.0
    min_resolution = 512
    
    # 오디오 길이
    audio = AudioSegment.from_file(audio_path)
    duration = len(audio) / 1000.0  # 초 단위
    
    # 오디오 길이에 따라 inference_steps 계산 (초당 약 12.5 프레임)
    # 최소 25 프레임, 최대 750 프레임 (60초 => 60*12.5=750)
    inference_steps = min(max(int(duration * 12.5), 25), 750)
    
    print(f"[INFO] Audio duration: {duration:.2f} seconds, using inference_steps: {inference_steps}")

    # 얼굴 인식 (face_info는 참고용)
    face_info = pipe.preprocess(img_path, expand_ratio=expand_ratio)
    print(f"[INFO] Face detection info: {face_info}")

    # 얼굴이 하나라도 검출되면(>0), 원본 이미지 비율 유지
    if face_info['face_num'] > 0:
        os.makedirs(os.path.dirname(res_video_path), exist_ok=True)
        
        # Sonic pipeline으로 비디오 생성
        pipe.process(
            img_path,
            audio_path,
            res_video_path,
            min_resolution=min_resolution,
            inference_steps=inference_steps,
            dynamic_scale=dynamic_scale
        )
        return res_video_path
    else:
        return -1


def process_sonic(image, audio, dynamic_scale):
    """
    Gradio 인터페이스 상에서 호출되는 함수.
    1. 이미지/오디오 입력 검증
    2. MD5 해시 통해 파일명 생성 후 캐싱
    3. 이미 결과 파일이 있으면 재활용, 없으면 새로 비디오 생성
    """
    if image is None:
        raise gr.Error("Please upload an image")
    if audio is None:
        raise gr.Error("Please upload an audio file")

    # 이미지 MD5 해시 계산
    buf_img = io.BytesIO()
    image.save(buf_img, format="PNG")
    img_bytes = buf_img.getvalue()
    img_md5 = get_md5(img_bytes)

    # 오디오 MD5 해시 계산
    sampling_rate, arr = audio[:2]
    if len(arr.shape) == 1:
        arr = arr[:, None]

    audio_segment = AudioSegment(
        arr.tobytes(),
        frame_rate=sampling_rate,
        sample_width=arr.dtype.itemsize,
        channels=arr.shape[1]
    )

    # (중요) Whisper 호환을 위해 mono/16kHz 변환
    audio_segment = audio_segment.set_channels(1)
    audio_segment = audio_segment.set_frame_rate(16000)

    # 최대 60초 제한
    MAX_DURATION_MS = 60000
    if len(audio_segment) > MAX_DURATION_MS:
        audio_segment = audio_segment[:MAX_DURATION_MS]

    buf_audio = io.BytesIO()
    audio_segment.export(buf_audio, format="wav")
    audio_bytes = buf_audio.getvalue()
    audio_md5 = get_md5(audio_bytes)

    # 파일 경로 생성
    image_path = os.path.abspath(os.path.join(tmp_path, f'{img_md5}.png'))
    audio_path = os.path.abspath(os.path.join(tmp_path, f'{audio_md5}.wav'))
    res_video_path = os.path.abspath(os.path.join(res_path, f'{img_md5}_{audio_md5}_{dynamic_scale}.mp4'))
    
    # 이미지/오디오 파일 캐싱
    if not os.path.exists(image_path):
        with open(image_path, "wb") as f:
            f.write(img_bytes)
    if not os.path.exists(audio_path):
        with open(audio_path, "wb") as f:
            f.write(audio_bytes)
    
    # 이미 결과가 존재하면 캐시된 결과 사용
    if os.path.exists(res_video_path):
        print(f"[INFO] Using cached result: {res_video_path}")
        return res_video_path
    else:
        print(f"[INFO] Generating new video with dynamic_scale={dynamic_scale}")
        video_result = get_video_res(image_path, audio_path, res_video_path, dynamic_scale)
        return video_result


def get_example():
    """예시 데이터를 로딩하는 더미 함수 (현재는 빈 리스트)."""
    return []


css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.main-header {
    text-align: center;
    color: #2a2a2a;
    margin-bottom: 2em;
}
.parameter-section {
    background-color: #f5f5f5;
    padding: 1em;
    border-radius: 8px;
    margin: 1em 0;
}
.example-section {
    margin-top: 2em;
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("""
        <div class="main-header">
            <h1>🎭 Sonic: Advanced Portrait Animation</h1>
            <p>Transform still images into dynamic videos synchronized with audio (up to 1 minute)</p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type='pil',
                label="Portrait Image",
                elem_id="image_input"
            )
            
            audio_input = gr.Audio(
                label="Voice/Audio Input (up to 1 minute)",
                elem_id="audio_input",
                type="numpy"
            )
            
            with gr.Column():
                dynamic_scale = gr.Slider(
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Animation Intensity",
                    info="Adjust to control movement intensity (0.5: subtle, 2.0: dramatic)"
                )
            
            process_btn = gr.Button(
                "Generate Animation",
                variant="primary",
                elem_id="process_btn"
            )
        
        with gr.Column():
            video_output = gr.Video(
                label="Generated Animation",
                elem_id="video_output"
            )
    
    process_btn.click(
        fn=process_sonic,
        inputs=[image_input, audio_input, dynamic_scale],
        outputs=video_output,
    )
    
    gr.Examples(
        examples=get_example(),
        fn=process_sonic,
        inputs=[image_input, audio_input, dynamic_scale],
        outputs=video_output,
        cache_examples=False
    )
    
    gr.HTML("""
        <div style="text-align: center; margin-top: 2em;">
            <div style="margin-bottom: 1em;">
                <a href="https://github.com/jixiaozhong/Sonic" target="_blank" style="text-decoration: none;">
                    <img src="https://img.shields.io/badge/GitHub-Repo-blue?style=for-the-badge&logo=github" alt="GitHub Repo">
                </a>
                <a href="https://arxiv.org/pdf/2411.16331" target="_blank" style="text-decoration: none;">
                    <img src="https://img.shields.io/badge/Paper-arXiv-red?style=for-the-badge&logo=arxiv" alt="arXiv Paper">
                </a>
            </div>
            <p>🔔 Note: For optimal results, use clear portrait images and high-quality audio (now supports up to 1 minute!)</p>
        </div>
    """)

# 공개 링크 생성
demo.launch(share=True)
import spaces
import gradio as gr
import os
import numpy as np
from pydub import AudioSegment
import hashlib
import io
from sonic import Sonic
from PIL import Image
import torch

# 초기 실행 시 필요한 모델들을 다운로드
cmd = (
    'python3 -m pip install "huggingface_hub[cli]"; '
    'huggingface-cli download LeonJoe13/Sonic --local-dir checkpoints; '
    'huggingface-cli download stabilityai/stable-video-diffusion-img2vid-xt --local-dir  checkpoints/stable-video-diffusion-img2vid-xt; '
    'huggingface-cli download openai/whisper-tiny --local-dir checkpoints/whisper-tiny;'
)
os.system(cmd)

pipe = Sonic()

def get_md5(content_bytes):
    """MD5 해시를 계산하여 32자리 문자열을 반환"""
    return hashlib.md5(content_bytes).hexdigest()

tmp_path = './tmp_path/'
res_path = './res_path/'
os.makedirs(tmp_path, exist_ok=True)
os.makedirs(res_path, exist_ok=True)


@spaces.GPU(duration=600)  # 긴 비디오 처리를 위해 duration 600초로 설정 (10분)
def get_video_res(img_path, audio_path, res_video_path, dynamic_scale=1.0):
    """
    Sonic pipeline으로부터 실제 비디오를 생성하는 함수.
    최대 60초 길이의 오디오에 대해 inference_steps를 결정하여,
    얼굴 탐지 후 영상 생성 작업을 수행함.
    """
    expand_ratio = 0.0
    min_resolution = 512
    
    # 오디오 길이
    audio = AudioSegment.from_file(audio_path)
    duration = len(audio) / 1000.0  # 초 단위
    
    # 오디오 길이에 따라 inference_steps 계산 (초당 약 12.5 프레임)
    # 최소 25 프레임, 최대 750 프레임 (60초 => 60*12.5=750)
    inference_steps = min(max(int(duration * 12.5), 25), 750)
    
    print(f"[INFO] Audio duration: {duration:.2f} seconds, using inference_steps: {inference_steps}")

    # 얼굴 인식 (face_info는 참고용)
    face_info = pipe.preprocess(img_path, expand_ratio=expand_ratio)
    print(f"[INFO] Face detection info: {face_info}")

    # 얼굴이 하나라도 검출되면(>0), 원본 이미지 비율 유지
    if face_info['face_num'] > 0:
        os.makedirs(os.path.dirname(res_video_path), exist_ok=True)
        
        # Sonic pipeline으로 비디오 생성
        pipe.process(
            img_path,
            audio_path,
            res_video_path,
            min_resolution=min_resolution,
            inference_steps=inference_steps,
            dynamic_scale=dynamic_scale
        )
        return res_video_path
    else:
        return -1


def process_sonic(image, audio, dynamic_scale):
    """
    Gradio 인터페이스 상에서 호출되는 함수.
    1. 이미지/오디오 입력 검증
    2. MD5 해시 통해 파일명 생성 후 캐싱
    3. 이미 결과 파일이 있으면 재활용, 없으면 새로 비디오 생성
    """
    if image is None:
        raise gr.Error("Please upload an image")
    if audio is None:
        raise gr.Error("Please upload an audio file")

    # 이미지 MD5 해시 계산
    buf_img = io.BytesIO()
    image.save(buf_img, format="PNG")
    img_bytes = buf_img.getvalue()
    img_md5 = get_md5(img_bytes)

    # 오디오 MD5 해시 계산
    sampling_rate, arr = audio[:2]
    if len(arr.shape) == 1:
        arr = arr[:, None]

    audio_segment = AudioSegment(
        arr.tobytes(),
        frame_rate=sampling_rate,
        sample_width=arr.dtype.itemsize,
        channels=arr.shape[1]
    )

    # (중요) Whisper 호환을 위해 mono/16kHz 변환
    audio_segment = audio_segment.set_channels(1)
    audio_segment = audio_segment.set_frame_rate(16000)

    # 최대 60초 제한
    MAX_DURATION_MS = 60000
    if len(audio_segment) > MAX_DURATION_MS:
        audio_segment = audio_segment[:MAX_DURATION_MS]

    buf_audio = io.BytesIO()
    audio_segment.export(buf_audio, format="wav")
    audio_bytes = buf_audio.getvalue()
    audio_md5 = get_md5(audio_bytes)

    # 파일 경로 생성
    image_path = os.path.abspath(os.path.join(tmp_path, f'{img_md5}.png'))
    audio_path = os.path.abspath(os.path.join(tmp_path, f'{audio_md5}.wav'))
    res_video_path = os.path.abspath(os.path.join(res_path, f'{img_md5}_{audio_md5}_{dynamic_scale}.mp4'))
    
    # 이미지/오디오 파일 캐싱
    if not os.path.exists(image_path):
        with open(image_path, "wb") as f:
            f.write(img_bytes)
    if not os.path.exists(audio_path):
        with open(audio_path, "wb") as f:
            f.write(audio_bytes)
    
    # 이미 결과가 존재하면 캐시된 결과 사용
    if os.path.exists(res_video_path):
        print(f"[INFO] Using cached result: {res_video_path}")
        return res_video_path
    else:
        print(f"[INFO] Generating new video with dynamic_scale={dynamic_scale}")
        video_result = get_video_res(image_path, audio_path, res_video_path, dynamic_scale)
        return video_result


def get_example():
    """예시 데이터를 로딩하는 더미 함수 (현재는 빈 리스트)."""
    return []


css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.main-header {
    text-align: center;
    color: #2a2a2a;
    margin-bottom: 2em;
}
.parameter-section {
    background-color: #f5f5f5;
    padding: 1em;
    border-radius: 8px;
    margin: 1em 0;
}
.example-section {
    margin-top: 2em;
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("""
        <div class="main-header">
            <h1>🎭 Sonic: Advanced Portrait Animation</h1>
            <p>Transform still images into dynamic videos synchronized with audio (up to 1 minute)</p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type='pil',
                label="Portrait Image",
                elem_id="image_input"
            )
            
            audio_input = gr.Audio(
                label="Voice/Audio Input (up to 1 minute)",
                elem_id="audio_input",
                type="numpy"
            )
            
            with gr.Column():
                dynamic_scale = gr.Slider(
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Animation Intensity",
                    info="Adjust to control movement intensity (0.5: subtle, 2.0: dramatic)"
                )
            
            process_btn = gr.Button(
                "Generate Animation",
                variant="primary",
                elem_id="process_btn"
            )
        
        with gr.Column():
            video_output = gr.Video(
                label="Generated Animation",
                elem_id="video_output"
            )
    
    process_btn.click(
        fn=process_sonic,
        inputs=[image_input, audio_input, dynamic_scale],
        outputs=video_output,
    )
    
    gr.Examples(
        examples=get_example(),
        fn=process_sonic,
        inputs=[image_input, audio_input, dynamic_scale],
        outputs=video_output,
        cache_examples=False
    )
    
    gr.HTML("""
        <div style="text-align: center; margin-top: 2em;">
            <div style="margin-bottom: 1em;">
                <a href="https://github.com/jixiaozhong/Sonic" target="_blank" style="text-decoration: none;">
                    <img src="https://img.shields.io/badge/GitHub-Repo-blue?style=for-the-badge&logo=github" alt="GitHub Repo">
                </a>
                <a href="https://arxiv.org/pdf/2411.16331" target="_blank" style="text-decoration: none;">
                    <img src="https://img.shields.io/badge/Paper-arXiv-red?style=for-the-badge&logo=arxiv" alt="arXiv Paper">
                </a>
            </div>
            <p>🔔 Note: For optimal results, use clear portrait images and high-quality audio (now supports up to 1 minute!)</p>
        </div>
    """)

# 공개 링크 생성
demo.launch(share=True)