Update sonic.py

sonic.py

@@ -1,4 +1,4 @@
-# sonic.py  (전체 파일)
+# sonic.py  ── 전체
 
 import os, math, glob, torch, cv2
 from PIL import Image
@@ -22,32 +22,29 @@ from src.dataset.face_align.align import AlignImage
 
 try:
     from safetensors.torch import load_file as safe_load
-except ImportError:       # safetensors 가 없으면 torch.load 만 사용
-    safe_load = None
+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",
+    """root 밑에서 keyword 가 포함된 .pth / .pt / .safetensors 하나 찾기"""
+    patterns = [f"**/*{keyword}*.pth",
+                f"**/*{keyword}*.pt",
                 f"**/*{keyword}*.safetensors"]
-    files = []
     for p in patterns:
-        files.extend(glob.glob(os.path.join(root, p), recursive=True))
-    return files[0] if files else None
+        files = glob.glob(os.path.join(root, p), recursive=True)
+        if files:
+            return files[0]
+    return None
 
 
-# -------------------------------------------------------------------
-#            single image + speech  →  video tensor
-# -------------------------------------------------------------------
+# --------------------------------------------------- speech → video
 def test(pipe, cfg, wav_enc, audio_pe, audio2bucket, image_encoder,
          width, height, batch):
 
-    # 배치 차원 맞추기
     for k, v in batch.items():
         if isinstance(v, torch.Tensor):
             batch[k] = v.unsqueeze(0).to(pipe.device).float()
@@ -59,17 +56,17 @@ def test(pipe, cfg, wav_enc, audio_pe, audio2bucket, image_encoder,
 
     audio_feature = batch["audio_feature"]          # (1,80,T)
     audio_len     = int(batch["audio_len"])
-    step          = max(1, int(cfg.step),)          # 최소 1
+    step          = max(1, int(cfg.step))
 
-    window = 16_000                                 # 1초 단위
+    window = 16_000
     audio_prompts, last_prompts = [], []
 
     for i in range(0, audio_feature.shape[-1], window):
         chunk = audio_feature[:, :, i:i+window]
-        hidden_layers = wav_enc.encoder(chunk, output_hidden_states=True).hidden_states
-        last_hidden   = wav_enc.encoder(chunk).last_hidden_state.unsqueeze(-2)
-        audio_prompts.append(torch.stack(hidden_layers, dim=2))
-        last_prompts.append(last_hidden)
+        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)
 
     if not audio_prompts:
         raise ValueError("[ERROR] No speech recognised in the provided audio.")
@@ -77,33 +74,29 @@ def test(pipe, cfg, wav_enc, audio_pe, audio2bucket, image_encoder,
     audio_prompts = torch.cat(audio_prompts, dim=1)
     last_prompts  = torch.cat(last_prompts , dim=1)
 
-    # padding 규칙
     audio_prompts = torch.cat(
         [torch.zeros_like(audio_prompts[:, :4]),
          audio_prompts,
-         torch.zeros_like(audio_prompts[:, :6])], dim=1)
-
+         torch.zeros_like(audio_prompts[:, :6])], 1)
     last_prompts = torch.cat(
         [torch.zeros_like(last_prompts[:, :24]),
          last_prompts,
-         torch.zeros_like(last_prompts[:, :26])], dim=1)
+         torch.zeros_like(last_prompts[:, :26])], 1)
 
-    total_tokens = audio_prompts.shape[1]
-    num_chunks   = max(1, math.ceil(total_tokens / (2*step)))
+    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], dim=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], dim=2)
+            bucket_clip = torch.cat([bucket_clip, pad], 2)
 
         motion = audio2bucket(bucket_clip, image_embeds) * 16 + 16
 
@@ -132,12 +125,10 @@ def test(pipe, cfg, wav_enc, audio_pe, audio2bucket, image_encoder,
         i2i_noise_strength=cfg.i2i_noise_strength,
     ).frames
 
-    return (video * 0.5 + 0.5).clamp(0, 1).unsqueeze(0).cpu()
+    return (video * .5 + .5).clamp(0,1).unsqueeze(0).cpu()
 
 
-# -------------------------------------------------------------------
-#                          Sonic  ✨
-# -------------------------------------------------------------------
+# ------------------------------------------------------------  Sonic
 class Sonic:
     config_file = os.path.join(BASE_DIR, "config/inference/sonic.yaml")
     config      = OmegaConf.load(config_file)
@@ -147,45 +138,40 @@ class Sonic:
         cfg.use_interframe = enable_interpolate_frame
         self.device        = f"cuda:{device_id}" if torch.cuda.is_available() and device_id >= 0 else "cpu"
 
-        # 가중치 루트
-        ckpt_root = os.path.join(BASE_DIR, "checkpoints", "Sonic")
-        cfg.pretrained_model_name_or_path = ckpt_root   # diffusers 형식
+        # 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")
 
-        self._load_models(cfg, ckpt_root)
+        self._load_models(cfg)
         print("Sonic init done")
 
-    # --------------------------------------------------------------
-    def _load_models(self, cfg, ckpt_root):
+    # --------------------------------------------- load all networks
+    def _load_models(self, cfg):
         dtype = {"fp16": torch.float16, "fp32": torch.float32, "bf16": torch.bfloat16}[cfg.weight_dtype]
 
-        # diffusers 기본 가중치
-        vae = AutoencoderKLTemporalDecoder.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="vae", variant="fp16")
-        sched = EulerDiscreteScheduler.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="scheduler")
-        image_enc = CLIPVisionModelWithProjection.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="image_encoder", variant="fp16")
-        unet = UNetSpatioTemporalConditionModel.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="unet", variant="fp16")
+        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")
         add_ip_adapters(unet, [32], [cfg.ip_audio_scale])
 
-        # ------------ 추가 체크포인트 (.pth / .safetensors) ------------
-        def _try_load(module, keyword):
-            path = _find_ckpt(ckpt_root, keyword)
+        def _load_extra(module, key):
+            path = _find_ckpt(self.ckpt_root, key)
             if not path:
-                print(f"[WARN] {keyword} checkpoint not found → skip")
+                print(f"[WARN] extra ckpt for '{key}' not found → skip")
                 return
-            print(f"[INFO] load {keyword} ckpt → {os.path.relpath(path, BASE_DIR)}")
-            if path.endswith(".safetensors") and safe_load is not None:
-                state = safe_load(path, device="cpu")
-            else:
-                state = torch.load(path, map_location="cpu")
+            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)
 
-        _try_load(unet,      "unet")
-        # audio adapters (필수)
         a2t = AudioProjModel(10, 5, 384, 1024, 1024, 32).to(self.device)
         a2b = Audio2bucketModel(50, 1, 384, 1024, 1024, 1, 2).to(self.device)
-        _try_load(a2t, "audio2token")
-        _try_load(a2b, "audio2bucket")
 
-        # whisper tiny
+        _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()
@@ -199,16 +185,16 @@ class Sonic:
             self.rife = RIFEModel(device=self.device)
             self.rife.load_model(os.path.join(BASE_DIR, "checkpoints/RIFE/"))
 
-        for m in (image_enc, vae, unet):
+        for m in (img_e, vae, unet):
             m.to(dtype)
 
-        self.pipe          = SonicPipeline(unet=unet, image_encoder=image_enc, vae=vae, scheduler=sched).to(device=self.device, dtype=dtype)
-        self.image_encoder = image_enc
+        self.pipe          = SonicPipeline(unet=unet, image_encoder=img_e, vae=vae, scheduler=sched).to(device=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):
         img = cv2.imread(image_path)
         h, w = img.shape[:2]
@@ -219,20 +205,19 @@ class Sonic:
                     "crop_bbox": process_bbox((x1, y1, x1+ww, y1+hh), expand_ratio, h, w)}
         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):
+                min_resolution=512, inference_steps=25,
+                dynamic_scale=1.0, keep_resolution=False, seed=None):
 
         cfg = self.config
         if seed is not None:
             cfg.seed = seed
-        cfg.num_inference_steps = inference_steps
-        cfg.motion_bucket_scale = dynamic_scale
+        cfg.num_inference_steps   = inference_steps
+        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,
@@ -252,7 +237,6 @@ class Sonic:
                      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):