Upload webgui-backup.py

webgui-backup.py

@@ -0,0 +1,311 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+webui
+'''
+
+import os
+import random
+from datetime import datetime
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+from diffusers import AutoencoderKL, DDIMScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+from src.models.unet_2d_condition import UNet2DConditionModel
+from src.models.unet_3d_echo import EchoUNet3DConditionModel
+from src.models.whisper.audio2feature import load_audio_model
+from src.pipelines.pipeline_echo_mimic import Audio2VideoPipeline
+from src.utils.util import save_videos_grid, crop_and_pad
+from src.models.face_locator import FaceLocator
+from moviepy.editor import VideoFileClip, AudioFileClip
+from facenet_pytorch import MTCNN
+import argparse
+
+import gradio as gr
+
+import huggingface_hub
+
+huggingface_hub.snapshot_download(
+    repo_id='BadToBest/EchoMimic',
+    local_dir='./pretrained_weights',
+    local_dir_use_symlinks=False,
+)
+
+# 환경 변수 대신 코드 내에서 직접 설정
+is_shared_ui = False  # 또는 True, 필요에 따라 설정
+
+# is_shared_ui의 값에 따라 available_property 설정
+available_property = not is_shared_ui
+
+# 이제 is_shared_ui와 available_property 변수는 코드 내에서 직접 관리됩니다.
+advanced_settings_label = "Advanced Settings"
+
+default_values = {
+    "width": 512,
+    "height": 512,
+    "length": 1200,
+    "seed": 420,
+    "facemask_dilation_ratio": 0.1,
+    "facecrop_dilation_ratio": 1.0,
+    "context_frames": 12,
+    "context_overlap": 3,
+    "cfg": 2.5,
+    "steps": 100,
+    "sample_rate": 16000,
+    "fps": 24,
+    "device": "cuda"
+}
+
+ffmpeg_path = os.getenv('FFMPEG_PATH')
+if ffmpeg_path is None:
+    print("please download ffmpeg-static and export to FFMPEG_PATH. \nFor example: export FFMPEG_PATH=/musetalk/ffmpeg-4.4-amd64-static")
+elif ffmpeg_path not in os.getenv('PATH'):
+    print("add ffmpeg to path")
+    os.environ["PATH"] = f"{ffmpeg_path}:{os.environ['PATH']}"
+
+
+config_path = "./configs/prompts/animation.yaml"
+config = OmegaConf.load(config_path)
+if config.weight_dtype == "fp16":
+    weight_dtype = torch.float16
+else:
+    weight_dtype = torch.float32
+
+device = "cuda"
+if not torch.cuda.is_available():
+    device = "cpu"
+
+inference_config_path = config.inference_config
+infer_config = OmegaConf.load(inference_config_path)
+
+############# model_init started #############
+## vae init
+vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path).to("cuda", dtype=weight_dtype)
+
+## reference net init
+reference_unet = UNet2DConditionModel.from_pretrained(
+    config.pretrained_base_model_path,
+    subfolder="unet",
+).to(dtype=weight_dtype, device=device)
+reference_unet.load_state_dict(torch.load(config.reference_unet_path, map_location="cpu"))
+
+## denoising net init
+if os.path.exists(config.motion_module_path):
+    ### stage1 + stage2
+    denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        config.motion_module_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(dtype=weight_dtype, device=device)
+else:
+    ### only stage1
+    denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        "",
+        subfolder="unet",
+        unet_additional_kwargs={
+            "use_motion_module": False,
+            "unet_use_temporal_attention": False,
+            "cross_attention_dim": infer_config.unet_additional_kwargs.cross_attention_dim
+        }
+    ).to(dtype=weight_dtype, device=device)
+
+denoising_unet.load_state_dict(torch.load(config.denoising_unet_path, map_location="cpu"), strict=False)
+
+## face locator init
+face_locator = FaceLocator(320, conditioning_channels=1, block_out_channels=(16, 32, 96, 256)).to(dtype=weight_dtype, device="cuda")
+face_locator.load_state_dict(torch.load(config.face_locator_path))
+
+## load audio processor params
+audio_processor = load_audio_model(model_path=config.audio_model_path, device=device)
+
+## load face detector params
+face_detector = MTCNN(image_size=320, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True, device=device)
+
+############# model_init finished #############
+
+sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
+scheduler = DDIMScheduler(**sched_kwargs)
+
+pipe = Audio2VideoPipeline(
+    vae=vae,
+    reference_unet=reference_unet,
+    denoising_unet=denoising_unet,
+    audio_guider=audio_processor,
+    face_locator=face_locator,
+    scheduler=scheduler,
+).to("cuda", dtype=weight_dtype)
+
+def select_face(det_bboxes, probs):
+    ## max face from faces that the prob is above 0.8
+    ## box: xyxy
+    if det_bboxes is None or probs is None:
+        return None
+    filtered_bboxes = []
+    for bbox_i in range(len(det_bboxes)):
+        if probs[bbox_i] > 0.8:
+            filtered_bboxes.append(det_bboxes[bbox_i])
+    if len(filtered_bboxes) == 0:
+        return None
+    sorted_bboxes = sorted(filtered_bboxes, key=lambda x:(x[3]-x[1]) * (x[2] - x[0]), reverse=True)
+    return sorted_bboxes[0]
+
+def process_video(uploaded_img, uploaded_audio, width, height, length, seed, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device):
+
+    if seed is not None and seed > -1:
+        generator = torch.manual_seed(seed)
+    else:
+        generator = torch.manual_seed(random.randint(100, 1000000))
+
+    #### face musk prepare
+    face_img = cv2.imread(uploaded_img)
+    face_mask = np.zeros((face_img.shape[0], face_img.shape[1])).astype('uint8')
+    det_bboxes, probs = face_detector.detect(face_img)
+    select_bbox = select_face(det_bboxes, probs)
+    if select_bbox is None:
+        face_mask[:, :] = 255
+    else:
+        xyxy = select_bbox[:4]
+        xyxy = np.round(xyxy).astype('int')
+        rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2]
+        r_pad = int((re - rb) * facemask_dilation_ratio)
+        c_pad = int((ce - cb) * facemask_dilation_ratio)
+        face_mask[rb - r_pad : re + r_pad, cb - c_pad : ce + c_pad] = 255
+        
+        #### face crop
+        r_pad_crop = int((re - rb) * facecrop_dilation_ratio)
+        c_pad_crop = int((ce - cb) * facecrop_dilation_ratio)
+        crop_rect = [max(0, cb - c_pad_crop), max(0, rb - r_pad_crop), min(ce + c_pad_crop, face_img.shape[1]), min(re + r_pad_crop, face_img.shape[0])]
+        face_img = crop_and_pad(face_img, crop_rect)
+        face_mask = crop_and_pad(face_mask, crop_rect)
+        face_img = cv2.resize(face_img, (width, height))
+        face_mask = cv2.resize(face_mask, (width, height))
+
+    ref_image_pil = Image.fromarray(face_img[:, :, [2, 1, 0]])
+    face_mask_tensor = torch.Tensor(face_mask).to(dtype=weight_dtype, device="cuda").unsqueeze(0).unsqueeze(0).unsqueeze(0) / 255.0
+    
+    video = pipe(
+        ref_image_pil,
+        uploaded_audio,
+        face_mask_tensor,
+        width,
+        height,
+        length,
+        steps,
+        cfg,
+        generator=generator,
+        audio_sample_rate=sample_rate,
+        context_frames=context_frames,
+        fps=fps,
+        context_overlap=context_overlap
+    ).videos
+
+    save_dir = Path("output/tmp")
+    save_dir.mkdir(exist_ok=True, parents=True)
+    output_video_path = save_dir / "output_video.mp4"
+    save_videos_grid(video, str(output_video_path), n_rows=1, fps=fps)
+
+    video_clip = VideoFileClip(str(output_video_path))
+    audio_clip = AudioFileClip(uploaded_audio)
+    final_output_path = save_dir / "output_video_with_audio.mp4"
+    video_clip = video_clip.set_audio(audio_clip)
+    video_clip.write_videofile(str(final_output_path), codec="libx264", audio_codec="aac")
+
+    return final_output_path
+  
+with gr.Blocks() as demo:
+    gr.Markdown('# Mimic FACE')
+
+    with gr.Row():
+        with gr.Column():
+            uploaded_img = gr.Image(type="filepath", label="Reference Image")
+            uploaded_audio = gr.Audio(type="filepath", label="Input Audio")
+            with gr.Accordion(label=advanced_settings_label, open=False):
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=128, maximum=1024, value=default_values["width"], interactive=available_property)
+                    height = gr.Slider(label="Height", minimum=128, maximum=1024, value=default_values["height"], interactive=available_property)
+                with gr.Row():
+                    length = gr.Slider(label="Length", minimum=100, maximum=5000, value=default_values["length"], interactive=available_property)
+                    seed = gr.Slider(label="Seed", minimum=0, maximum=10000, value=default_values["seed"], interactive=available_property)
+                with gr.Row():
+                    facemask_dilation_ratio = gr.Slider(label="Facemask Dilation Ratio", minimum=0.0, maximum=1.0, step=0.01, value=default_values["facemask_dilation_ratio"], interactive=available_property)
+                    facecrop_dilation_ratio = gr.Slider(label="Facecrop Dilation Ratio", minimum=0.0, maximum=1.0, step=0.01, value=default_values["facecrop_dilation_ratio"], interactive=available_property)
+                with gr.Row():
+                    context_frames = gr.Slider(label="Context Frames", minimum=0, maximum=50, step=1, value=default_values["context_frames"], interactive=available_property)
+                    context_overlap = gr.Slider(label="Context Overlap", minimum=0, maximum=10, step=1, value=default_values["context_overlap"], interactive=available_property)
+                with gr.Row():
+                    cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0, step=0.1, value=default_values["cfg"], interactive=available_property)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=default_values["steps"], interactive=available_property)
+                with gr.Row():
+                    sample_rate = gr.Slider(label="Sample Rate", minimum=8000, maximum=48000, step=1000, value=default_values["sample_rate"], interactive=available_property)
+                    fps = gr.Slider(label="FPS", minimum=1, maximum=60, step=1, value=default_values["fps"], interactive=available_property)
+                    device = gr.Radio(label="Device", choices=["cuda", "cpu"], value=default_values["device"], interactive=available_property)
+            generate_button = gr.Button("Generate Video")
+        with gr.Column():
+            output_video = gr.Video()
+            gr.Examples(
+                label = "Portrait examples",
+                examples = [
+                    ['assets/test_imgs/a.png'],
+                ],
+                inputs = [uploaded_img]
+            )
+            gr.Examples(
+                label = "Audio examples",
+                examples = [
+                    ['assets/test_audios/chunnuanhuakai.wav'],
+                ],
+                inputs = [uploaded_audio]
+            )
+
+    def generate_video(uploaded_img, uploaded_audio, width, height, length, seed, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device):
+
+        final_output_path = process_video(
+            uploaded_img, uploaded_audio, width, height, length, seed, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device
+        )        
+        output_video= final_output_path
+        return final_output_path
+
+    generate_button.click(
+        generate_video,
+        inputs=[
+            uploaded_img,
+            uploaded_audio,
+            width,
+            height,
+            length,
+            seed,
+            facemask_dilation_ratio,
+            facecrop_dilation_ratio,
+            context_frames,
+            context_overlap,
+            cfg,
+            steps,
+            sample_rate,
+            fps,
+            device
+        ],
+        outputs=output_video,
+        api_name="generate_video_api"  # Expose API endpoint
+    )
+
+
+parser = argparse.ArgumentParser(description='Mimic FACE')
+parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Server name')
+parser.add_argument('--server_port', type=int, default=7860, help='Server port')
+args = parser.parse_args()
+
+
+
+if __name__ == '__main__':
+#    demo.launch(
+    demo.queue(max_size=4).launch(   
+        server_name=args.server_name,
+        server_port=args.server_port,
+        show_api=True  # Enable API documentation
+    )