Update sonic.py

sonic.py

@@ -1,9 +1,15 @@
-# sonic.py  ── 전체
-
-import os, math, glob, torch, cv2
+# sonic.py
+# ---------------------------------------------------------------------
+#  Sonic – single-image + speech → talking-head video  (offline edition)
+# ---------------------------------------------------------------------
+import os, math
+from typing import Dict, Any, List
+
+import torch
 from PIL import Image
 from omegaconf import OmegaConf
 from tqdm import tqdm
+import cv2
 
 from diffusers import AutoencoderKLTemporalDecoder
 from diffusers.schedulers import EulerDiscreteScheduler
@@ -12,7 +18,8 @@ from transformers import WhisperModel, CLIPVisionModelWithProjection, AutoFeatur
 from src.utils.util import save_videos_grid, seed_everything
 from src.dataset.test_preprocess import process_bbox, image_audio_to_tensor
 from src.models.base.unet_spatio_temporal_condition import (
-    UNetSpatioTemporalConditionModel, add_ip_adapters,
+    UNetSpatioTemporalConditionModel,
+    add_ip_adapters,
 )
 from src.pipelines.pipeline_sonic import SonicPipeline
 from src.models.audio_adapter.audio_proj import AudioProjModel
@@ -20,115 +27,165 @@ from src.models.audio_adapter.audio_to_bucket import Audio2bucketModel
 from src.utils.RIFE.RIFE_HDv3 import RIFEModel
 from src.dataset.face_align.align import AlignImage
 
-try:
-    from safetensors.torch import load_file as safe_load
-except ImportError:
-    safe_load = None                                           # safetensors 미설치 시 대비
-
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
-# ------------------------------------------------------------ utils
-def _find_ckpt(root: str, keyword: str):
-    """root 밑에서 keyword 가 포함된 .pth / .pt / .safetensors 하나 찾기"""
-    patterns = [f"**/*{keyword}*.pth",
-                f"**/*{keyword}*.pt",
-                f"**/*{keyword}*.safetensors"]
-    for p in patterns:
-        files = glob.glob(os.path.join(root, p), recursive=True)
-        if files:
-            return files[0]
-    return None
-
-
-# --------------------------------------------------- speech → video
-def test(pipe, cfg, wav_enc, audio_pe, audio2bucket, image_encoder,
-         width, height, batch):
-
+# ------------------------------------------------------------------ #
+#                     헬퍼 : diffusers 경로 자동 찾기                  #
+# ------------------------------------------------------------------ #
+def _locate_diffusers_dir(root: str) -> str:
+    """
+    `root` 하위 디렉터리에서 diffusers 스냅샷(model_index.json or config.json)
+    이 들어 있는 실제 모델 폴더를 찾아서 반환한다. 존재하지 않으면 오류.
+    """
+    for cur, _dirs, files in os.walk(root):
+        if {"model_index.json", "config.json"} & set(files):
+            return cur
+    raise FileNotFoundError(
+        f"[ERROR] No diffusers model files found under '{root}'. "
+        "Check that the checkpoint was downloaded correctly."
+    )
+
+
+# ------------------------------------------------------------------ #
+#                        영상 생성용 내부 함수                         #
+# ------------------------------------------------------------------ #
+def _gen_video_tensor(
+    pipe: SonicPipeline,
+    cfg: OmegaConf,
+    wav_enc: WhisperModel,
+    audio_pe: AudioProjModel,
+    audio2bucket: Audio2bucketModel,
+    image_encoder: CLIPVisionModelWithProjection,
+    width: int,
+    height: int,
+    batch: Dict[str, torch.Tensor],
+) -> torch.Tensor:
+    """
+    single 이미지 + 오디오 feature → video tensor (C,T,H,W)
+    """
+
+    # -------- batch 차원 보정 --------------------------------------
     for k, v in batch.items():
         if isinstance(v, torch.Tensor):
             batch[k] = v.unsqueeze(0).to(pipe.device).float()
 
-    ref_img   = batch["ref_img"]
+    ref_img   = batch["ref_img"]        # (1,C,H,W)
     clip_img  = batch["clip_images"]
     face_mask = batch["face_mask"]
     image_embeds = image_encoder(clip_img).image_embeds
 
-    audio_feature = batch["audio_feature"]          # (1,80,T)
-    audio_len     = int(batch["audio_len"])
-    step          = max(1, int(cfg.step))
+    audio_feat: torch.Tensor = batch["audio_feature"]          # (1, 80, T)
+    audio_len:  int          = int(batch["audio_len"])         # scalar
+    step:       int          = int(cfg.step)
 
-    window = 16_000
-    audio_prompts, last_prompts = [], []
+    # step 이 전체 길이보다 크면 최소 1 로 보정
+    if audio_len < step:
+        step = max(1, audio_len)
 
-    for i in range(0, audio_feature.shape[-1], window):
-        chunk = audio_feature[:, :, i:i+window]
-        hidden = wav_enc.encoder(chunk, output_hidden_states=True).hidden_states
-        last   = wav_enc.encoder(chunk).last_hidden_state.unsqueeze(-2)
-        audio_prompts.append(torch.stack(hidden, dim=2))
-        last_prompts.append(last)
+    # -------- Whisper encoder 1초 단위로 수행 ----------------------
+    window = 16_000                                             # 1-s chunk
+    aud_prompts: List[torch.Tensor] = []
+    last_prompts: List[torch.Tensor] = []
 
-    if not audio_prompts:
+    for i in range(0, audio_feat.shape[-1], window):
+        chunk = audio_feat[:, :, i : i + window]
+
+        # 모든 hidden-states / 마지막 hidden-state
+        layers: List[torch.Tensor] = wav_enc.encoder(
+            chunk, output_hidden_states=True
+        ).hidden_states
+        last_hidden = wav_enc.encoder(chunk).last_hidden_state  # (1, 80, 384)
+
+        # Whisper layer 는 6개 → AudioProj 가 기대하는 5개로 truncate
+        prompt = torch.stack(layers, dim=2)[:, :, :5]           # (1,80,5,384)
+        aud_prompts.append(prompt)
+        last_prompts.append(last_hidden.unsqueeze(-2))          # (1,80,1,384)
+
+    if len(aud_prompts) == 0:
         raise ValueError("[ERROR] No speech recognised in the provided audio.")
 
-    audio_prompts = torch.cat(audio_prompts, dim=1)
-    last_prompts  = torch.cat(last_prompts , dim=1)
+    # concat 뒤 padding 규칙 적용
+    aud_prompts = torch.cat(aud_prompts, dim=1)                 # (1, 80*…, 5, 384)
+    last_prompts = torch.cat(last_prompts, dim=1)               # (1, 80*…, 1, 384)
 
-    audio_prompts = torch.cat(
-        [torch.zeros_like(audio_prompts[:, :4]),
-         audio_prompts,
-         torch.zeros_like(audio_prompts[:, :6])], 1)
+    aud_prompts = torch.cat(
+        [torch.zeros_like(aud_prompts[:, :4]), aud_prompts, torch.zeros_like(aud_prompts[:, :6])],
+        dim=1,
+    )
     last_prompts = torch.cat(
-        [torch.zeros_like(last_prompts[:, :24]),
-         last_prompts,
-         torch.zeros_like(last_prompts[:, :26])], 1)
-
-    num_chunks = max(1, math.ceil(audio_prompts.shape[1] / (2*step)))
-
-    ref_list, audio_list, uncond_list, buckets = [], [], [], []
-    for i in tqdm(range(num_chunks)):
-        st = i * 2 * step
-        cond = audio_prompts[:, st: st+10]
-        if cond.shape[2] < 10:
-            pad = torch.zeros_like(cond[:, :, :10-cond.shape[2]])
-            cond = torch.cat([cond, pad], 2)
-
-        bucket_clip = last_prompts[:, st: st+50]
-        if bucket_clip.shape[2] < 50:
-            pad = torch.zeros_like(bucket_clip[:, :, :50-bucket_clip.shape[2]])
-            bucket_clip = torch.cat([bucket_clip, pad], 2)
-
-        motion = audio2bucket(bucket_clip, image_embeds) * 16 + 16
-
-        ref_list.append(ref_img[0])
-        audio_list.append(audio_pe(cond).squeeze(0))
-        uncond_list.append(audio_pe(torch.zeros_like(cond)).squeeze(0))
-        buckets.append(motion[0])
-
-    video = pipe(
-        ref_img, clip_img, face_mask,
-        audio_list, uncond_list, buckets,
-        height=height, width=width,
-        num_frames=len(audio_list),
-        decode_chunk_size=cfg.decode_chunk_size,
-        motion_bucket_scale=cfg.motion_bucket_scale,
-        fps=cfg.fps,
-        noise_aug_strength=cfg.noise_aug_strength,
-        min_guidance_scale1=cfg.min_appearance_guidance_scale,
-        max_guidance_scale1=cfg.max_appearance_guidance_scale,
-        min_guidance_scale2=cfg.audio_guidance_scale,
-        max_guidance_scale2=cfg.audio_guidance_scale,
-        overlap=cfg.overlap,
-        shift_offset=cfg.shift_offset,
-        frames_per_batch=cfg.n_sample_frames,
-        num_inference_steps=cfg.num_inference_steps,
-        i2i_noise_strength=cfg.i2i_noise_strength,
-    ).frames
-
-    return (video * .5 + .5).clamp(0,1).unsqueeze(0).cpu()
-
-
-# ------------------------------------------------------------  Sonic
+        [torch.zeros_like(last_prompts[:, :24]), last_prompts, torch.zeros_like(last_prompts[:, :26])],
+        dim=1,
+    )
+
+    # --------  f=10 / w=5 로 clip 자르기 --------------------------
+    ref_list, aud_list, uncond_list, mb_list = [], [], [], []
+
+    total_tokens = aud_prompts.shape[1]
+    n_chunks = max(1, math.ceil(total_tokens / (2 * step)))
+
+    for i in tqdm(range(n_chunks), desc="audio-chunks", ncols=0):
+        s = i * 2 * step
+
+        cond_clip = aud_prompts[:, s : s + 10]                  # (1,10,5,384)
+        if cond_clip.shape[1] < 10:                             # 뒤쪽 padding
+            pad = torch.zeros_like(cond_clip[:, : 10 - cond_clip.shape[1]])
+            cond_clip = torch.cat([cond_clip, pad], dim=1)
+
+        bucket_clip = last_prompts[:, s : s + 50]               # (1,50,1,384)
+        if bucket_clip.shape[1] < 50:
+            pad = torch.zeros_like(bucket_clip[:, : 50 - bucket_clip.shape[1]])
+            bucket_clip = torch.cat([bucket_clip, pad], dim=1)
+
+        # (bz,f,w,b,c) 5-D 로 변환
+        cond_clip      = cond_clip.unsqueeze(3)                 # (1,10,5,1,384)
+        bucket_clip    = bucket_clip.unsqueeze(3)               # (1,50,1,1,384)
+        uncond_clip    = torch.zeros_like(cond_clip)
+
+        motion_bucket  = audio2bucket(bucket_clip, image_embeds) * 16 + 16
+
+        ref_list      .append(ref_img[0])
+        aud_list      .append(audio_pe(cond_clip).squeeze(0)[0])      # (ctx,1024)
+        uncond_list   .append(audio_pe(uncond_clip).squeeze(0)[0])    # (ctx,1024)
+        mb_list       .append(motion_bucket[0])
+
+    # --------  UNet 파이프라인 실행 --------------------------------
+    video = (
+        pipe(
+            ref_img,
+            clip_img,
+            face_mask,
+            aud_list,
+            uncond_list,
+            mb_list,
+            height=height,
+            width=width,
+            num_frames=len(aud_list),
+            decode_chunk_size=cfg.decode_chunk_size,
+            motion_bucket_scale=cfg.motion_bucket_scale,
+            fps=cfg.fps,
+            noise_aug_strength=cfg.noise_aug_strength,
+            min_guidance_scale1=cfg.min_appearance_guidance_scale,
+            max_guidance_scale1=cfg.max_appearance_guidance_scale,
+            min_guidance_scale2=cfg.audio_guidance_scale,
+            max_guidance_scale2=cfg.audio_guidance_scale,
+            overlap=cfg.overlap,
+            shift_offset=cfg.shift_offset,
+            frames_per_batch=cfg.n_sample_frames,
+            num_inference_steps=cfg.num_inference_steps,
+            i2i_noise_strength=cfg.i2i_noise_strength,
+        ).frames
+        * 0.5
+        + 0.5
+    ).clamp(0, 1)
+
+    # (B,C,T,H,W)   → (C,T,H,W)
+    return video.to(pipe.device).squeeze(0).cpu()
+
+
+# ------------------------------------------------------------------ #
+#                         Sonic  –  main class                        #
+# ------------------------------------------------------------------ #
 class Sonic:
     config_file = os.path.join(BASE_DIR, "config/inference/sonic.yaml")
     config      = OmegaConf.load(config_file)
@@ -136,96 +193,87 @@ class Sonic:
     def __init__(self, device_id: int = 0, enable_interpolate_frame: bool = True):
         cfg                = self.config
         cfg.use_interframe = enable_interpolate_frame
-        self.device        = f"cuda:{device_id}" if torch.cuda.is_available() and device_id >= 0 else "cpu"
 
-        # diffusers 베이스 모델은 ⇣ (config.json 포함)
-        self.diffusers_root = os.path.join(BASE_DIR, "checkpoints", "stable-video-diffusion-img2vid-xt")
-        # 추가 pth/pt/safetensors 는 ⇣
-        self.ckpt_root      = os.path.join(BASE_DIR, "checkpoints", "Sonic")
+        # diffusers 모델 상위 폴더 (로컬 다운로드 경로)
+        self.diffusers_root = os.path.join(BASE_DIR, cfg.pretrained_model_name_or_path)
+        self.device = (
+            f"cuda:{device_id}" if device_id >= 0 and torch.cuda.is_available() else "cpu"
+        )
 
         self._load_models(cfg)
         print("Sonic init done")
 
-    def _locate_diffusers_dir(root: str) -> str:
-        """
-        root 아래에서 model_index.json 또는 config.json 이 존재하는
-        디렉터리를 찾아 반환. (snapshots/<sha>/ … 형식 대응)
-        """
-        for cur, _dirs, files in os.walk(root):
-            if {"model_index.json", "config.json"} & set(files):
-                return cur
-        raise FileNotFoundError(
-            f"[ERROR] diffusers model files(model_index.json/config.json) "
-            f"not found under {root}"
-        )
-
-    
-    # --------------------------------------------- load all networks
+    # -------------------------------------------------------------- #
     def _load_models(self, cfg):
+        # dtype
         dtype = {"fp16": torch.float16, "fp32": torch.float32, "bf16": torch.bfloat16}[cfg.weight_dtype]
-        diff_root = _locate_diffusers_dir(self.diffusers_root)   # ★★ 핵심 추가
 
-        vae   = AutoencoderKLTemporalDecoder.from_pretrained(self.diffusers_root, subfolder="vae", variant="fp16")
-        sched = EulerDiscreteScheduler.from_pretrained(self.diffusers_root, subfolder="scheduler")
-        img_e = CLIPVisionModelWithProjection.from_pretrained(self.diffusers_root, subfolder="image_encoder", variant="fp16")
-        unet  = UNetSpatioTemporalConditionModel.from_pretrained(self.diffusers_root, subfolder="unet", variant="fp16")
+        diff_root = _locate_diffusers_dir(self.diffusers_root)
+
+        # diffusers 모듈들
+        vae   = AutoencoderKLTemporalDecoder.from_pretrained(diff_root, subfolder="vae", variant="fp16")
+        sched = EulerDiscreteScheduler.from_pretrained(diff_root, subfolder="scheduler")
+        img_e = CLIPVisionModelWithProjection.from_pretrained(diff_root, subfolder="image_encoder", variant="fp16")
+        unet  = UNetSpatioTemporalConditionModel.from_pretrained(diff_root, subfolder="unet", variant="fp16")
         add_ip_adapters(unet, [32], [cfg.ip_audio_scale])
 
-        def _load_extra(module, key):
-            path = _find_ckpt(self.ckpt_root, key)
-            if not path:
-                print(f"[WARN] extra ckpt for '{key}' not found → skip")
-                return
-            print(f"[INFO] load {key} → {os.path.relpath(path, BASE_DIR)}")
-            state = safe_load(path, device="cpu") if (safe_load and path.endswith(".safetensors")) else torch.load(path, map_location="cpu")
-            module.load_state_dict(state, strict=False)
-
-        a2t = AudioProjModel(10, 5, 384, 1024, 1024, 32).to(self.device)
-        a2b = Audio2bucketModel(50, 1, 384, 1024, 1024, 1, 2).to(self.device)
-
-        _load_extra(unet, "unet")
-        _load_extra(a2t,  "audio2token")
-        _load_extra(a2b,  "audio2bucket")
-
-        whisper = WhisperModel.from_pretrained(
-            os.path.join(BASE_DIR, "checkpoints/whisper-tiny")
-        ).to(self.device).eval()
+        # 오디오 어댑터
+        a2t = AudioProjModel(seq_len=10, blocks=5, channels=384,
+                             intermediate_dim=1024, output_dim=1024, context_tokens=32).to(self.device)
+        a2b = Audio2bucketModel(seq_len=50, blocks=1, channels=384,
+                                clip_channels=1024, intermediate_dim=1024, output_dim=1,
+                                context_tokens=2).to(self.device)
+
+        # 체크포인트 로드
+        a2t.load_state_dict(torch.load(os.path.join(BASE_DIR, cfg.audio2token_checkpoint_path), map_location="cpu"))
+        a2b.load_state_dict(torch.load(os.path.join(BASE_DIR, cfg.audio2bucket_checkpoint_path), map_location="cpu"))
+        unet.load_state_dict(torch.load(os.path.join(BASE_DIR, cfg.unet_checkpoint_path), map_location="cpu"))
+
+        # Whisper
+        whisper = WhisperModel.from_pretrained(os.path.join(BASE_DIR, "checkpoints/whisper-tiny")).to(self.device).eval()
         whisper.requires_grad_(False)
 
-        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
-            os.path.join(BASE_DIR, "checkpoints/whisper-tiny")
-        )
+        # 이미지 / 얼굴 / 보간
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(os.path.join(BASE_DIR, "checkpoints/whisper-tiny"))
         self.face_det = AlignImage(self.device, det_path=os.path.join(BASE_DIR, "checkpoints/yoloface_v5m.pt"))
         if cfg.use_interframe:
             self.rife = RIFEModel(device=self.device)
             self.rife.load_model(os.path.join(BASE_DIR, "checkpoints/RIFE/"))
 
-        for m in (img_e, vae, unet):
+        # dtype 적용
+        for m in (vae, img_e, unet):
             m.to(dtype)
 
-        self.pipe          = SonicPipeline(unet=unet, image_encoder=img_e, vae=vae, scheduler=sched).to(device=self.device, dtype=dtype)
+        self.pipe          = SonicPipeline(unet=unet, image_encoder=img_e, vae=vae, scheduler=sched).to(self.device, dtype=dtype)
         self.image_encoder = img_e
         self.audio2token   = a2t
         self.audio2bucket  = a2b
         self.whisper       = whisper
 
-    # --------------------------------------------- preprocess helpers
-    def preprocess(self, image_path: str, expand_ratio: float = 1.0):
+    # -------------------------------------------------------------- #
+    def preprocess(self, image_path: str, expand_ratio: float = 1.0) -> Dict[str, Any]:
         img = cv2.imread(image_path)
         h, w = img.shape[:2]
         _, _, bboxes = self.face_det(img, maxface=True)
         if bboxes:
             x1, y1, ww, hh = bboxes[0]
-            return {"face_num": 1,
-                    "crop_bbox": process_bbox((x1, y1, x1+ww, y1+hh), expand_ratio, h, w)}
+            crop = process_bbox((x1, y1, x1 + ww, y1 + hh), expand_ratio, h, w)
+            return {"face_num": 1, "crop_bbox": crop}
         return {"face_num": 0, "crop_bbox": None}
 
-    # --------------------------------------------------------------- run
+    # -------------------------------------------------------------- #
     @torch.no_grad()
-    def process(self, image_path, audio_path, output_path,
-                min_resolution=512, inference_steps=25,
-                dynamic_scale=1.0, keep_resolution=False, seed=None):
-
+    def process(
+        self,
+        image_path: str,
+        audio_path: str,
+        output_path: str,
+        min_resolution: int = 512,
+        inference_steps: int = 25,
+        dynamic_scale: float = 1.0,
+        keep_resolution: bool = False,
+        seed: int | None = None,
+    ) -> int:
         cfg = self.config
         if seed is not None:
             cfg.seed = seed
@@ -233,10 +281,15 @@ class Sonic:
         cfg.motion_bucket_scale   = dynamic_scale
         seed_everything(cfg.seed)
 
+        # 이미지·오디오 tensor 변환
         data = image_audio_to_tensor(
-            self.face_det, self.feature_extractor,
-            image_path, audio_path,
-            limit=-1, image_size=min_resolution, area=cfg.area
+            self.face_det,
+            self.feature_extractor,
+            image_path,
+            audio_path,
+            limit=-1,
+            image_size=min_resolution,
+            area=cfg.area,
         )
         if data is None:
             return -1
@@ -244,27 +297,31 @@ class Sonic:
         h, w = data["ref_img"].shape[-2:]
         if keep_resolution:
             im = Image.open(image_path)
-            resolution = f"{im.width//2*2}x{im.height//2*2}"
+            resolution = f"{(im.width // 2) * 2}x{(im.height // 2) * 2}"
         else:
             resolution = f"{w}x{h}"
 
-        video = test(self.pipe, cfg, self.whisper, self.audio2token,
-                     self.audio2bucket, self.image_encoder,
-                     w, h, data)
+        # video tensor 생성
+        video = _gen_video_tensor(
+            self.pipe, cfg, self.whisper, self.audio2token, self.audio2bucket,
+            self.image_encoder, w, h, data,
+        )
 
+        # 중간 프레임 보간
         if cfg.use_interframe:
-            out, frames = video.to(self.device), []
-            for i in tqdm(range(out.shape[2]-1), ncols=0):
-                mid = self.rife.inference(out[:,:,i], out[:,:,i+1]).clamp(0,1).detach()
-                frames.extend([out[:,:,i], mid])
-            frames.append(out[:,:,-1])
-            video = torch.stack(frames, 2).cpu()
-
+            out = video.to(self.device)
+            frames = []
+            for i in tqdm(range(out.shape[1] - 1), desc="interpolate", ncols=0):
+                frames.extend([out[:, i], self.rife.inference(out[:, i], out[:, i + 1]).clamp(0, 1)])
+            frames.append(out[:, -1])
+            video = torch.stack(frames, 1).cpu()  # (C,T',H,W)
+
+        # 저장
         tmp = output_path.replace(".mp4", "_noaudio.mp4")
-        save_videos_grid(video, tmp, n_rows=video.shape[0],
-                         fps=cfg.fps*(2 if cfg.use_interframe else 1))
+        save_videos_grid(video.unsqueeze(0), tmp, n_rows=1, fps=cfg.fps * (2 if cfg.use_interframe else 1))
         os.system(
             f"ffmpeg -loglevel error -y -i '{tmp}' -i '{audio_path}' -s {resolution} "
-            f"-vcodec libx264 -acodec aac -crf 18 -shortest '{output_path}'")
+            f"-vcodec libx264 -acodec aac -crf 18 -shortest '{output_path}'"
+        )
         os.remove(tmp)
         return 0