File size: 13,586 Bytes
7b4dc6f
 
 
 
 
 
 
 
79d88c4
 
0329637
7b4dc6f
0329637
 
 
79d88c4
0329637
79d88c4
c260fe0
0329637
7b4dc6f
 
0329637
c260fe0
79d88c4
 
 
 
0329637
79d88c4
 
914dc02
7b4dc6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79d88c4
 
0329637
85ad908
7b4dc6f
0329637
 
 
3c16281
7b4dc6f
 
 
1fb410d
7b4dc6f
 
 
f40c908
7b4dc6f
 
 
 
0178f77
7b4dc6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0329637
0178f77
7b4dc6f
 
 
0178f77
7b4dc6f
 
 
 
0329637
7b4dc6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1fb410d
 
 
 
c260fe0
f40c908
85ad908
c260fe0
7b4dc6f
 
 
 
 
1fb410d
a47303a
0329637
c260fe0
7b4dc6f
a47303a
7b4dc6f
85ad908
1fb410d
7b4dc6f
 
 
 
 
 
 
1fb410d
79d88c4
7b4dc6f
 
 
 
 
 
 
 
 
 
 
 
 
 
0329637
 
7b4dc6f
 
1fb410d
 
0329637
 
 
7b4dc6f
 
0329637
1fb410d
7b4dc6f
a47303a
0329637
 
 
1fb410d
7b4dc6f
 
0329637
 
 
 
 
7b4dc6f
 
0329637
79d88c4
7b4dc6f
612b064
7b4dc6f
 
 
 
 
 
 
 
 
 
 
1fb410d
0329637
 
a47303a
 
1fb410d
 
7b4dc6f
0329637
7b4dc6f
 
 
 
 
 
 
6ee08fc
0329637
 
1fb410d
0329637
 
 
7b4dc6f
0329637
 
1fb410d
7b4dc6f
 
 
 
 
5d6304b
7b4dc6f
0329637
7b4dc6f
 
 
 
 
 
 
 
0329637
7b4dc6f
0329637
 
7b4dc6f
 
0329637
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
# 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
from transformers import WhisperModel, CLIPVisionModelWithProjection, AutoFeatureExtractor

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,
)
from src.pipelines.pipeline_sonic import SonicPipeline
from src.models.audio_adapter.audio_proj import AudioProjModel
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

BASE_DIR = os.path.dirname(os.path.abspath(__file__))


# ------------------------------------------------------------------ #
#                     헬퍼 : 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"]        # (1,C,H,W)
    clip_img  = batch["clip_images"]
    face_mask = batch["face_mask"]
    image_embeds = image_encoder(clip_img).image_embeds

    audio_feat: torch.Tensor = batch["audio_feature"]          # (1, 80, T)
    audio_len:  int          = int(batch["audio_len"])         # scalar
    step:       int          = int(cfg.step)

    # step 이 전체 길이보다 크면 최소 1 로 보정
    if audio_len < step:
        step = max(1, audio_len)

    # -------- Whisper encoder 1초 단위로 수행 ----------------------
    window = 16_000                                             # 1-s chunk
    aud_prompts: List[torch.Tensor] = []
    last_prompts: List[torch.Tensor] = []

    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.")

    # 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)

    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])],
        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)

    def __init__(self, device_id: int = 0, enable_interpolate_frame: bool = True):
        cfg                = self.config
        cfg.use_interframe = enable_interpolate_frame

        # 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 _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)

        # 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])

        # 오디오 어댑터
        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.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/"))

        # 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(self.device, dtype=dtype)
        self.image_encoder = img_e
        self.audio2token   = a2t
        self.audio2bucket  = a2b
        self.whisper       = whisper

    # -------------------------------------------------------------- #
    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]
            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}

    # -------------------------------------------------------------- #
    @torch.no_grad()
    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
        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,
            limit=-1,
            image_size=min_resolution,
            area=cfg.area,
        )
        if data is None:
            return -1

        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}"
        else:
            resolution = f"{w}x{h}"

        # 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 = 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.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}'"
        )
        os.remove(tmp)
        return 0