Update app.py

app.py

@@ -10,6 +10,14 @@ import torch
 from diffusers import AutoencoderKL, TCDScheduler
 from diffusers.models.model_loading_utils import load_state_dict
 from huggingface_hub import hf_hub_download
+from pathlib import Path
+import torchaudio
+from einops import rearrange
+from scipy.io import wavfile
+from transformers import pipeline
+
+# 환경 변수 설정으로 torch.load 체크 우회 (임시 해결책)
+os.environ["TRANSFORMERS_ALLOW_UNSAFE_DESERIALIZATION"] = "1"
 
 # Spaces GPU
 try:
@@ -23,22 +31,19 @@ except:
                 return func
             return decorator
 
-# 환경 변수 설정
-os.environ["TRANSFORMERS_ALLOW_UNSAFE_DESERIALIZATION"] = "1"
-
-# MMAudio 관련 임포트
+# MMAudio imports
 try:
     import mmaudio
-    from mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate, load_video, make_video,
-                                    setup_eval_logging)
-    from mmaudio.model.flow_matching import FlowMatching
-    from mmaudio.model.networks import MMAudio, get_my_mmaudio
-    from mmaudio.model.sequence_config import SequenceConfig
-    from mmaudio.model.utils.features_utils import FeaturesUtils
-    MMAUDIO_AVAILABLE = True
 except ImportError:
-    MMAUDIO_AVAILABLE = False
-    logging.warning("MMAudio not available. Sound generation will be disabled.")
+    os.system("pip install -e .")
+    import mmaudio
+
+from mmaudio.eval_utils import (ModelConfig, all_model_cfg, generate, load_video, make_video,
+                                setup_eval_logging)
+from mmaudio.model.flow_matching import FlowMatching
+from mmaudio.model.networks import MMAudio, get_my_mmaudio
+from mmaudio.model.sequence_config import SequenceConfig
+from mmaudio.model.utils.features_utils import FeaturesUtils
 
 # ControlNet 모델 로드
 try:
@@ -89,52 +94,59 @@ except Exception as e:
     logging.error(f"Failed to load outpainting models: {str(e)}")
     OUTPAINT_MODEL_LOADED = False
 
-# MMAudio 모델 설정 및 로드
-if MMAUDIO_AVAILABLE:
+# MMAudio 모델 설정
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+else:
+    device = torch.device("cpu")
+
+dtype = torch.bfloat16
+
+# MMAudio 모델 초기화
+try:
+    model_mmaudio: ModelConfig = all_model_cfg['large_44k_v2']
+    model_mmaudio.download_if_needed()
+    output_dir = Path('./output/gradio')
+    setup_eval_logging()
+    
+    # 번역기 설정
     try:
-        # CUDA 설정
-        if torch.cuda.is_available():
-            device = torch.device("cuda")
-            torch.backends.cuda.matmul.allow_tf32 = True
-            torch.backends.cudnn.allow_tf32 = True
-            torch.backends.cudnn.benchmark = True
-        else:
-            device = torch.device("cpu")
-        
-        dtype = torch.bfloat16
-        
-        # 모델 설정
-        model_cfg: ModelConfig = all_model_cfg['large_44k_v2']
-        model_cfg.download_if_needed()
-        
-        setup_eval_logging()
-        
-        # 모델 로드
-        def get_mmaudio_model():
-            with torch.cuda.device(device):
-                seq_cfg = model_cfg.seq_cfg
-                net: MMAudio = get_my_mmaudio(model_cfg.model_name).to(device, dtype).eval()
-                net.load_weights(torch.load(model_cfg.model_path, map_location=device, weights_only=True))
-                logging.info(f'Loaded MMAudio weights from {model_cfg.model_path}')
-                
-                feature_utils = FeaturesUtils(
-                    tod_vae_ckpt=model_cfg.vae_path,
-                    synchformer_ckpt=model_cfg.synchformer_ckpt,
-                    enable_conditions=True,
-                    mode=model_cfg.mode,
-                    bigvgan_vocoder_ckpt=model_cfg.bigvgan_16k_path,
-                    need_vae_encoder=False
-                ).to(device, dtype).eval()
-                
-                return net, feature_utils, seq_cfg
-        
-        mmaudio_net, mmaudio_feature_utils, mmaudio_seq_cfg = get_mmaudio_model()
-        MMAUDIO_LOADED = True
+        translator = pipeline("translation", 
+                             model="Helsinki-NLP/opus-mt-ko-en", 
+                             device="cpu",
+                             use_fast=True,
+                             trust_remote_code=False)
     except Exception as e:
-        logging.error(f"Failed to load MMAudio models: {str(e)}")
-        MMAUDIO_LOADED = False
-else:
-    MMAUDIO_LOADED = False
+        logging.warning(f"Failed to load translation model: {e}")
+        translator = None
+    
+    def get_mmaudio_model() -> tuple[MMAudio, FeaturesUtils, SequenceConfig]:
+        with torch.cuda.device(device):
+            seq_cfg = model_mmaudio.seq_cfg
+            net: MMAudio = get_my_mmaudio(model_mmaudio.model_name).to(device, dtype).eval()
+            net.load_weights(torch.load(model_mmaudio.model_path, map_location=device, weights_only=True))
+            logging.info(f'Loaded weights from {model_mmaudio.model_path}')
+
+            feature_utils = FeaturesUtils(
+                tod_vae_ckpt=model_mmaudio.vae_path,
+                synchformer_ckpt=model_mmaudio.synchformer_ckpt,
+                enable_conditions=True,
+                mode=model_mmaudio.mode,
+                bigvgan_vocoder_ckpt=model_mmaudio.bigvgan_16k_path,
+                need_vae_encoder=False
+            ).to(device, dtype).eval()
+
+            return net, feature_utils, seq_cfg
+
+    net_mmaudio, feature_utils, seq_cfg = get_mmaudio_model()
+    MMAUDIO_MODEL_LOADED = True
+except Exception as e:
+    logging.error(f"Failed to load MMAudio models: {str(e)}")
+    MMAUDIO_MODEL_LOADED = False
+    translator = None
 
 # API URLs
 TEXT2IMG_API_URL = "http://211.233.58.201:7896"
@@ -192,52 +204,6 @@ def generate_text_to_image(prompt, width, height, guidance, inference_steps, see
         logging.error(f"Image generation error: {str(e)}")
         return None, f"오류: {str(e)}"
 
-@spaces.GPU(duration=60)
-@torch.inference_mode()
-def video_to_audio(video_path, prompt, negative_prompt="music", seed=0, num_steps=25, cfg_strength=4.5, target_duration=8.0):
-    """비디오에 사운드를 추가하는 함수"""
-    if not MMAUDIO_LOADED:
-        logging.error("MMAudio model not loaded")
-        return video_path
-    
-    try:
-        rng = torch.Generator(device=device)
-        rng.manual_seed(seed)
-        fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
-        
-        # 비디오 로드 - target_duration 사용
-        clip_frames, sync_frames, actual_duration = load_video(video_path, target_duration)
-        clip_frames = clip_frames.unsqueeze(0)
-        sync_frames = sync_frames.unsqueeze(0)
-        mmaudio_seq_cfg.duration = actual_duration
-        mmaudio_net.update_seq_lengths(mmaudio_seq_cfg.latent_seq_len, mmaudio_seq_cfg.clip_seq_len, mmaudio_seq_cfg.sync_seq_len)
-        
-        # 오디오 생성
-        audios = generate(clip_frames,
-                          sync_frames, [prompt],
-                          negative_text=[negative_prompt],
-                          feature_utils=mmaudio_feature_utils,
-                          net=mmaudio_net,
-                          fm=fm,
-                          rng=rng,
-                          cfg_strength=cfg_strength)
-        audio = audios.float().cpu()[0]
-        
-        # 비디오와 오디오 결합
-        video_save_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
-        make_video(video_path,
-                   video_save_path,
-                   audio,
-                   sampling_rate=mmaudio_seq_cfg.sampling_rate,
-                   duration_sec=mmaudio_seq_cfg.duration)
-        
-        return video_save_path
-    except Exception as e:
-        logging.error(f"Video to audio error: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return video_path
-
 def generate_video_from_image(image, prompt="", length=4.0):
     if image is None:
         return None
@@ -270,25 +236,6 @@ def generate_video_from_image(image, prompt="", length=4.0):
         logging.error(f"Video generation error: {str(e)}")
         return None
 
-def add_sound_to_video(video_path, sound_prompt, sound_negative_prompt="music"):
-    if not video_path or not MMAUDIO_LOADED:
-        return video_path
-    
-    try:
-        return video_to_audio(
-            video_path=video_path,
-            prompt=sound_prompt,
-            negative_prompt=sound_negative_prompt,
-            seed=random.randint(0, 9999999),
-            num_steps=25,
-            cfg_strength=4.5,
-            target_duration=8.0  # 기본값 사용
-        )
-    except Exception as e:
-        logging.error(f"Sound addition error: {str(e)}")
-        return video_path
-
-
 def prepare_image_and_mask(image, width, height, overlap_percentage, alignment):
     """이미지와 마스크를 준비하는 함수"""
     if image is None:
@@ -339,7 +286,7 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, alignment):
     mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
     
-    # 마스크 영역 그리기
+    # 마스크 영역 그리기 (영어 정렬과 매칭)
     white_gaps_patch = 2
     
     left_overlap = margin_x + overlap_x if alignment != "왼쪽" else margin_x
@@ -354,6 +301,27 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, alignment):
     
     return background, mask
 
+def preview_outpaint(image, width, height, overlap_percentage, alignment):
+    """아웃페인팅 미리보기"""
+    background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, alignment)
+    if background is None:
+        return None
+    
+    # 미리보기 이미지 생성
+    preview = background.copy().convert('RGBA')
+    
+    # 반투명 빨간색 오버레이
+    red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))
+    
+    # 마스크 적용
+    red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
+    red_mask.paste(red_overlay, (0, 0), mask)
+    
+    # 오버레이 합성
+    preview = Image.alpha_composite(preview, red_mask)
+    
+    return preview
+
 @spaces.GPU(duration=24)
 def outpaint_image(image, prompt, width, height, overlap_percentage, alignment, num_steps=8):
     """이미지 아웃페인팅 실행"""
@@ -410,6 +378,59 @@ def outpaint_image(image, prompt, width, height, overlap_percentage, alignment,
         logging.error(f"Outpainting error: {str(e)}")
         return background if 'background' in locals() else None
 
+# MMAudio 관련 함수들
+def translate_prompt(text):
+    try:
+        if translator is None:
+            return text
+            
+        if text and any(ord(char) >= 0x3131 and ord(char) <= 0xD7A3 for char in text):
+            with torch.no_grad():
+                translation = translator(text)[0]['translation_text']
+            return translation
+        return text
+    except Exception as e:
+        logging.error(f"Translation error: {e}")
+        return text
+
+@spaces.GPU
+@torch.inference_mode()
+def video_to_audio(video: gr.Video, prompt: str, negative_prompt: str, seed: int, num_steps: int,
+                   cfg_strength: float, duration: float):
+    if not MMAUDIO_MODEL_LOADED:
+        return None
+    
+    prompt = translate_prompt(prompt)
+    negative_prompt = translate_prompt(negative_prompt)
+
+    rng = torch.Generator(device=device)
+    rng.manual_seed(seed)
+    fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=num_steps)
+
+    clip_frames, sync_frames, duration = load_video(video, duration)
+    clip_frames = clip_frames.unsqueeze(0)
+    sync_frames = sync_frames.unsqueeze(0)
+    seq_cfg.duration = duration
+    net_mmaudio.update_seq_lengths(seq_cfg.latent_seq_len, seq_cfg.clip_seq_len, seq_cfg.sync_seq_len)
+
+    audios = generate(clip_frames,
+                      sync_frames, [prompt],
+                      negative_text=[negative_prompt],
+                      feature_utils=feature_utils,
+                      net=net_mmaudio,
+                      fm=fm,
+                      rng=rng,
+                      cfg_strength=cfg_strength)
+    audio = audios.float().cpu()[0]
+
+    video_save_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+    make_video(video,
+               video_save_path,
+               audio,
+               sampling_rate=seq_cfg.sampling_rate,
+               duration_sec=seq_cfg.duration)
+    return video_save_path
+
 # CSS
 css = """
 :root {
@@ -435,7 +456,7 @@ css = """
     padding: 20px !important;
     margin-bottom: 20px !important;
 }
-#generate-btn, #video-btn, #outpaint-btn {
+#generate-btn, #video-btn, #outpaint-btn, #preview-btn, #audio-btn {
     background: linear-gradient(135deg, #ff9a9e, #fad0c4) !important;
     font-size: 1.1rem !important;
     padding: 12px 24px !important;
@@ -448,7 +469,7 @@ css = """
 """
 
 # Gradio Interface
-demo = gr.Blocks(css=css, title="AI 이미지 & 비디오 생성기")
+demo = gr.Blocks(css=css, title="AI 이미지 & 비디오 & 오디오 생성기")
 
 with demo:
     gr.Markdown("# 🎨 Ginigen 스튜디오")
@@ -504,27 +525,6 @@ with demo:
                             info="1초에서 60초까지 선택 가능합니다"
                         )
                         
-                        # 사운드 생성 옵션 추가
-                        sound_generation = gr.Radio(
-                            choices=["사운드 없음", "사운드 생성"],
-                            value="사운드 없음",
-                            label="사운드 옵션",
-                            info="비디오에 사운드를 추가할지 선택하세요"
-                        )
-                        
-                        # 사운드 관련 입력 필드 (조건부 표시)
-                        with gr.Column(visible=False) as sound_options:
-                            sound_prompt = gr.Textbox(
-                                label="사운드 프롬프트 (선택)",
-                                placeholder="생성할 사운드를 설명하세요... (비워두면 비디오 프롬프트 사용)",
-                                lines=2
-                            )
-                            sound_negative_prompt = gr.Textbox(
-                                label="사운드 네거티브 프롬프트",
-                                value="music",
-                                lines=1
-                            )
-                        
                         video_btn = gr.Button("🎬 비디오로 변환", variant="secondary", elem_id="video-btn")
                 
                 # 출력 컬럼
@@ -590,6 +590,7 @@ with demo:
                             label="추론 스텝"
                         )
                         
+                        preview_btn = gr.Button("👁️ 미리보기", elem_id="preview-btn")
                         outpaint_btn = gr.Button("🎨 아웃페인팅 실행", variant="primary", elem_id="outpaint-btn")
                 
                 # 출력 컬럼
@@ -597,7 +598,60 @@ with demo:
                     with gr.Group(elem_classes="panel-box"):
                         gr.Markdown("### 🖼️ 결과")
                         
+                        preview_image = gr.Image(label="미리보기")
                         outpaint_result = gr.Image(label="아웃페인팅 결과")
+        
+        # 세 번째 탭: 비디오 + 오디오
+        with gr.Tab("비디오 + 오디오", elem_classes="tabitem"):
+            with gr.Row(equal_height=True):
+                # 입력 컬럼
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🎥 비디오 업로드")
+                        
+                        audio_video_input = gr.Video(
+                            label="입력 비디오",
+                            sources=["upload"]
+                        )
+                        
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🎵 오디오 생성 설정")
+                        
+                        audio_prompt = gr.Textbox(
+                            label="프롬프트 (한글 지원)" if MMAUDIO_MODEL_LOADED and translator else "프롬프트",
+                            placeholder="생성하고 싶은 오디오를 설명하세요... (예: 평화로운 피아노 음악)",
+                            lines=3
+                        )
+                        
+                        audio_negative_prompt = gr.Textbox(
+                            label="네거티브 프롬프트",
+                            value="music",
+                            placeholder="원하지 않는 요소...",
+                            lines=2
+                        )
+                        
+                        with gr.Row():
+                            audio_seed = gr.Number(label="시드", value=0)
+                            audio_steps = gr.Number(label="스텝", value=25)
+                        
+                        with gr.Row():
+                            audio_cfg = gr.Number(label="가이던스 스케일", value=4.5)
+                            audio_duration = gr.Number(label="지속시간 (초)", value=8)
+                        
+                        audio_btn = gr.Button("🎵 오디오 생성 및 합성", variant="primary", elem_id="audio-btn")
+                
+                # 출력 컬럼
+                with gr.Column(scale=1):
+                    with gr.Group(elem_classes="panel-box"):
+                        gr.Markdown("### 🎬 생성 결과")
+                        
+                        output_video_with_audio = gr.Video(
+                            label="오디오가 추가된 비디오",
+                            interactive=False
+                        )
+                        
+                        if not MMAUDIO_MODEL_LOADED:
+                            gr.Markdown("⚠️ MMAudio 모델을 로드하지 못했습니다. 이 기능은 사용할 수 없습니다.")
     
     # 이벤트 연결 - 첫 번째 탭
     size_preset.change(update_dimensions, [size_preset], [width, height])
@@ -608,39 +662,32 @@ with demo:
         [output_image, output_seed]
     )
     
-    # 사운드 옵션 표시/숨김
-    def toggle_sound_options(choice):
-        return gr.update(visible=(choice == "사운드 생성"))
-    
-    sound_generation.change(
-        toggle_sound_options,
-        [sound_generation],
-        [sound_options]
-    )
-
-
     video_btn.click(
-        generate_video_from_image,
-        [output_image, video_prompt, video_length],  # 원래대로 3개 매개변수만
-        [output_video]
-    )
-
-# 사운드 추가는 별도 버튼으로
-    sound_btn = gr.Button("🔊 비디오에 사운드 추가", visible=False)
-    sound_btn.click(
-        add_sound_to_video,
-        [output_video, sound_prompt, sound_negative_prompt],
+        lambda img, v_prompt, length: generate_video_from_image(img, v_prompt, length) if img is not None else None,
+        [output_image, video_prompt, video_length],
         [output_video]
     )
-
     
     # 이벤트 연결 - 두 번째 탭
     outpaint_size_preset.change(update_dimensions, [outpaint_size_preset], [outpaint_width, outpaint_height])
     
+    preview_btn.click(
+        preview_outpaint,
+        [input_image, outpaint_width, outpaint_height, overlap_percentage, alignment],
+        [preview_image]
+    )
+    
     outpaint_btn.click(
         outpaint_image,
         [input_image, outpaint_prompt, outpaint_width, outpaint_height, overlap_percentage, alignment, outpaint_steps],
         [outpaint_result]
     )
+    
+    # 이벤트 연결 - 세 번째 탭
+    audio_btn.click(
+        video_to_audio,
+        [audio_video_input, audio_prompt, audio_negative_prompt, audio_seed, audio_steps, audio_cfg, audio_duration],
+        [output_video_with_audio]
+    )
 
 demo.launch()