Update app.py

app.py

@@ -1,126 +1,138 @@
-import os, io, hashlib, gradio as gr, spaces
+# app.py
+import os, io, hashlib, spaces, gradio as gr
 from pydub import AudioSegment
-from PIL import Image
-from sonic import Sonic
-
-# --------------------------------------------------------------
-# 1) 필요 체크포인트(두 개)만 받는다
-# --------------------------------------------------------------
-DL_CMDS = [
-    # Sonic(가중치, unet 등)
-    "huggingface-cli download LeonJoe13/Sonic --local-dir checkpoints/LeonJoe13-Sonic --local-dir-use-symlinks False -q",
-    # stable-video-diffusion-img2vid-xt (VAE/UNet/CLIP)
-    "huggingface-cli download stabilityai/stable-video-diffusion-img2vid-xt "
-    "--local-dir checkpoints/stable-video-diffusion-img2vid-xt --local-dir-use-symlinks False -q",
-    # whisper-tiny
-    "huggingface-cli download openai/whisper-tiny --local-dir checkpoints/whisper-tiny --local-dir-use-symlinks False -q",
-]
-for cmd in DL_CMDS:
-    os.system(cmd)
-
-pipe = Sonic()   # 위에서 모델 경로를 자동으로 찾음
+from PIL   import Image
+import numpy as np
+from sonic import Sonic          # <-- 수정된 sonic.py 사용
 
 # ------------------------------------------------------------------
-def md5(b: bytes) -> str:          # 빠른 32-byte 해시
+# 1. 필요한 모델·라이브러리 설치 & 체크포인트 다운로드
+# ------------------------------------------------------------------
+SETUP_CMD = (
+    # 필수 라이브러리
+    'python -m pip install --quiet "huggingface_hub[cli]" accelerate pydub Pillow '
+    # Sonic 체크포인트 → checkpoints/Sonic/*
+    '&& huggingface-cli download LeonJoe13/Sonic '
+    '--local-dir checkpoints/Sonic --local-dir-use-symlinks False '
+    # Whisper-tiny (음성 인코더)
+    '&& huggingface-cli download openai/whisper-tiny '
+    '--local-dir checkpoints/whisper-tiny --local-dir-use-symlinks False '
+)
+os.system(SETUP_CMD)
+
+# ------------------------------------------------------------------
+# 2. 파이프라인 초기화 (GPU 한 번만)
+# ------------------------------------------------------------------
+pipe = Sonic()     # 오류가 사라진 상태로 초기화
+
+# ------------------------------------------------------------------
+# 3. 유틸리티
+# ------------------------------------------------------------------
+def _md5(b: bytes) -> str:
     return hashlib.md5(b).hexdigest()
 
-TMP_DIR, RES_DIR = "./tmp_path", "./res_path"
-os.makedirs(TMP_DIR, exist_ok=True)
-os.makedirs(RES_DIR, exist_ok=True)
+TMP_DIR  = "tmp_path"
+RES_DIR  = "res_path"
+os.makedirs(TMP_DIR,  exist_ok=True)
+os.makedirs(RES_DIR,  exist_ok=True)
 
 # ------------------------------------------------------------------
-@spaces.GPU(duration=600)          # 최대 10분까지 GPU 사용
-def get_video_res(img_p: str, wav_p: str, out_p: str, scale: float):
-    """실제 Sonic 파이프라인 호출(얼굴 체크·프레임·interpolate 포함)"""
-    audio = AudioSegment.from_file(wav_p)
-    dur   = len(audio) / 1000.0
-    steps = max(25, min(int(dur * 12.5), 750))   # 12.5fps 기준
-
-    print(f"[INFO] Audio duration {dur:.2f}s ➜ steps {steps}")
-
-    face = pipe.preprocess(img_p)
-    print("[INFO] Face detection:", face)
-    if face["face_num"] == 0:
-        return -1                                 # 얼굴 없음
-
-    pipe.process(img_p, wav_p, out_p,
-                 min_resolution=512,
-                 inference_steps=steps,
-                 dynamic_scale=scale)
-    return out_p
+# 4. 실제 비디오 생성 (GPU 태스크)
+# ------------------------------------------------------------------
+@spaces.GPU(duration=600)     # 최대 10분
+def _render_video(img_path: str,
+                  audio_path: str,
+                  out_path: str,
+                  dynamic_scale: float = 1.0) -> str | int:
+
+    min_resolution = 512
+    audio          = AudioSegment.from_file(audio_path)
+    duration_sec   = len(audio) / 1000.0
+    steps          = int(np.clip(duration_sec * 12.5, 25, 750))
+
+    print(f"[INFO] Audio duration={duration_sec:.2f}s → inference_steps={steps}")
+
+    face_info = pipe.preprocess(img_path)
+    print(f"[INFO] Face detection info: {face_info}")
+
+    if face_info["face_num"] == 0:
+        return -1        # 얼굴 미검출
+
+    os.makedirs(os.path.dirname(out_path), exist_ok=True)
+    pipe.process(
+        img_path, audio_path, out_path,
+        min_resolution=min_resolution,
+        inference_steps=steps,
+        dynamic_scale=dynamic_scale,
+    )
+    return out_path
 
 # ------------------------------------------------------------------
-def run_sonic(image, audio, scale):
-    """Gradio 버튼 연결 함수 (캐싱·전처리)"""
-    if image is None: raise gr.Error("Please upload an image.")
-    if audio is None: raise gr.Error("Please upload an audio file.")
-
-    # ---- 이미지 저장 & 해시 -------------------------------------------------
-    buf_img = io.BytesIO(); image.save(buf_img, "PNG")
-    img_key = md5(buf_img.getvalue())
-    img_path = os.path.join(TMP_DIR, f"{img_key}.png")
+# 5. Gradio 래퍼
+# ------------------------------------------------------------------
+def run_sonic(image, audio, dynamic_scale):
+
+    if image is None:
+        raise gr.Error("Please upload an image.")
+    if audio is None:
+        raise gr.Error("Please upload an audio file.")
+
+    # ── 이미지 캐시 ───────────────────────────────────────────────
+    buf_i = io.BytesIO(); image.save(buf_i, format="PNG")
+    img_hash = _md5(buf_i.getvalue())
+    img_path = os.path.join(TMP_DIR, f"{img_hash}.png")
     if not os.path.exists(img_path):
-        with open(img_path, "wb") as f: f.write(buf_img.getvalue())
-
-    # ---- 오디오 → mono/16kHz WAV (≤60 s) -----------------------------------
-    sr, arr = audio[:2]
-    arr = arr if arr.ndim == 2 else arr[:, None]
-    seg = AudioSegment(arr.tobytes(), frame_rate=sr,
-                       sample_width=arr.dtype.itemsize,
-                       channels=arr.shape[1]
-           ).set_channels(1).set_frame_rate(16_000)[:60_000]
-    buf_wav = io.BytesIO(); seg.export(buf_wav, format="wav")
-    wav_key = md5(buf_wav.getvalue())
-    wav_path = os.path.join(TMP_DIR, f"{wav_key}.wav")
-    if not os.path.exists(wav_path):
-        with open(wav_path, "wb") as f: f.write(buf_wav.getvalue())
-
-    # ---- 결과 파일 경로 -----------------------------------------------------
-    out_path = os.path.join(RES_DIR, f"{img_key}_{wav_key}_{scale}.mp4")
-
-    # ---- 캐시 확인 ---------------------------------------------------------
+        with open(img_path, "wb") as f: f.write(buf_i.getvalue())
+
+    # ── 오디오 캐시 (mono/16 kHz, ≤60 s) ─────────────────────────
+    rate, arr = audio[:2]
+    if arr.ndim == 1: arr = arr[:, None]
+    seg = AudioSegment(arr.tobytes(), frame_rate=rate,
+                       sample_width=arr.dtype.itemsize, channels=arr.shape[1])
+    seg = seg.set_channels(1).set_frame_rate(16000)[:60_000]
+    buf_a = io.BytesIO(); seg.export(buf_a, format="wav")
+    aud_hash = _md5(buf_a.getvalue())
+    aud_path = os.path.join(TMP_DIR, f"{aud_hash}.wav")
+    if not os.path.exists(aud_path):
+        with open(aud_path, "wb") as f: f.write(buf_a.getvalue())
+
+    # ── 결과 경로 ────────────────────────────────────────────────
+    out_path = os.path.join(
+        RES_DIR, f"{img_hash}_{aud_hash}_{dynamic_scale:.1f}.mp4"
+    )
+
     if os.path.exists(out_path):
-        print("[INFO] Cached video used.")
+        print(f"[INFO] Cache hit → {out_path}")
         return out_path
 
-    print(f"[INFO] Generating video (scale={scale}) …")
-    res = get_video_res(img_path, wav_path, out_path, scale)
-    if res == -1:
-        raise gr.Error("No face detected in the image.")
-    return res
+    print(f"[INFO] Generating video (dynamic_scale={dynamic_scale}) …")
+    return _render_video(img_path, aud_path, out_path, dynamic_scale)
 
 # ------------------------------------------------------------------
-#                     Gradio UI
+# 6. Gradio UI
 # ------------------------------------------------------------------
-css = """
-.gradio-container {font-family: Arial, sans-serif;}
-.main-header     {text-align:center;color:#2a2a2a;margin-bottom:2em;}
+CSS = """
+.gradio-container{font-family:Arial, sans-serif}
+.main-header{text-align:center;color:#2a2a2a;margin-bottom:2em}
 """
 
-with gr.Blocks(css=css) as demo:
-    gr.HTML(
-        "<div class='main-header'>"
-        "<h1>🎭 Sonic: Portrait-to-Video Animator</h1>"
-        "<p>Create talking-head videos (≤60 s audio)</p>"
-        "</div>"
-    )
+with gr.Blocks(css=CSS) as demo:
+    gr.HTML("""
+    <div class="main-header">
+        <h1>🎭 Sonic - Portrait Animation</h1>
+        <p>Turn a single photo into a talking-head video (≤1 min audio)</p>
+    </div>""")
 
     with gr.Row():
         with gr.Column():
-            img_in = gr.Image(type="pil", label="Portrait Image")
-            aud_in = gr.Audio(type="numpy", label="Voice/Audio (≤60 s)")
-            scale  = gr.Slider(0.5, 2.0, 1.0, 0.1,
-                               label="Animation Intensity")
-            btn    = gr.Button("Generate", variant="primary")
-        vid_out = gr.Video(label="Generated Animation")
-
-    btn.click(run_sonic, inputs=[img_in, aud_in, scale], outputs=vid_out)
-
-    gr.HTML(
-        "<div style='text-align:center;margin-top:1.5em'>"
-        "<a href='https://github.com/jixiaozhong/Sonic' target='_blank'>GitHub</a> | "
-        "<a href='https://arxiv.org/pdf/2411.16331' target='_blank'>Paper</a>"
-        "</div>"
-    )
+            img_in  = gr.Image(type="pil",  label="Portrait Image")
+            aud_in  = gr.Audio(label="Voice / Audio (≤60 s)", type="numpy")
+            scale   = gr.Slider(0.5, 2.0, 1.0, step=0.1,
+                                label="Animation Intensity")
+            btn     = gr.Button("Generate Animation", variant="primary")
+        with gr.Column():
+            vid_out = gr.Video(label="Result")
+
+    btn.click(run_sonic, [img_in, aud_in, scale], vid_out)
 
 demo.launch(share=True)