Update app.py

app.py

@@ -1,48 +1,19 @@
-"""
-InfiniteTalk - Talking Video Generator
-Gradio Space for HuggingFace
-"""
-
 import os
-import sys
 import random
 import logging
-import warnings
-from pathlib import Path
-
-import gradio as gr
 import torch
-import numpy as np
-import librosa
-import soundfile as sf
-import pyloudnorm as pyln
+import gradio as gr
 from PIL import Image
-from einops import rearrange
-
-# Import utilities
 from utils.model_loader import ModelManager
 from utils.gpu_manager import gpu_manager
 import wan
-from wan.configs import SIZE_CONFIGS, WAN_CONFIGS
 from wan.utils.utils import cache_image, cache_video, is_video
 from wan.utils.multitalk_utils import save_video_ffmpeg
-from transformers import Wav2Vec2FeatureExtractor
-from src.audio_analysis.wav2vec2 import Wav2Vec2Model
-
-# Set environment variables before importing Torch
-os.environ["TORCHVISION_DISABLE_META_REGISTRATIONS"] = "1"
-os.environ["TORCH_LOGS"] = "-all"  # Reduce torch logging noise
-
-# Suppress warnings
-warnings.filterwarnings('ignore')
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Add current directory to path
-sys.path.insert(0, str(Path(__file__).parent))
-
 # Global variables
 model_manager = None
 models_loaded = False
@@ -73,61 +44,6 @@ def initialize_models(progress=gr.Progress()):
         logger.error(f"Error initializing models: {e}")
         raise gr.Error(f"Failed to initialize models: {str(e)}")
 
-
-def loudness_norm(audio_array, sr=16000, lufs=-20.0):
-    """Normalize audio loudness using pyloudnorm"""
-    try:
-        meter = pyln.Meter(sr)
-        loudness = meter.integrated_loudness(audio_array)
-        if abs(loudness) > 100:  # Skip if loudness measurement failed
-            return audio_array
-        normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
-        return normalized_audio
-    except Exception as e:
-        logger.warning(f"Loudness normalization failed: {e}, returning original audio")
-        return audio_array
-
-
-def process_audio(audio_path, target_sr=16000):
-    """Process audio file for InfiniteTalk"""
-    try:
-        # Load audio with librosa
-        audio, sr = librosa.load(audio_path, sr=target_sr)
-
-        # Normalize loudness
-        audio = loudness_norm(audio, sr)
-
-        # Ensure mono
-        if len(audio.shape) > 1:
-            audio = np.mean(audio, axis=1)
-
-        return audio, sr
-
-    except Exception as e:
-        logger.error(f"Error processing audio: {e}")
-        raise gr.Error(f"Audio processing failed: {str(e)}")
-
-
-def validate_inputs(image_or_video, audio, resolution, steps):
-    """Validate user inputs"""
-    errors = []
-
-    if image_or_video is None:
-        errors.append("Please upload an image or video")
-
-    if audio is None:
-        errors.append("Please upload an audio file")
-
-    if resolution not in ["480p", "720p"]:
-        errors.append("Invalid resolution selected")
-
-    if not (20 <= steps <= 50):
-        errors.append("Steps must be between 20 and 50")
-
-    if errors:
-        raise gr.Error(" | ".join(errors))
-
-
 def generate_video(
     image_or_video,
     audio_file,
@@ -139,7 +55,6 @@ def generate_video(
 ):
     """Generate talking video from image or dub existing video"""
     try:
-        # Check if GPU is available
         if not torch.cuda.is_available():
             raise gr.Error(
                 "⚠️ GPU not available. This Space requires GPU hardware to generate videos."
@@ -149,32 +64,17 @@ def generate_video(
         if not models_loaded:
             initialize_models(progress)
 
-        # Validate inputs
-        validate_inputs(image_or_video, audio_file, resolution, steps)
-
-        # GPU memory check
-        gpu_manager.print_memory_usage("Initial - ")
-
         progress(0.1, desc="Processing audio...")
 
-        # Process audio
-        audio, sr = process_audio(audio_file)
-        audio_duration = len(audio) / sr
-        logger.info(f"Audio duration: {audio_duration:.2f}s")
+        # Process audio (add your audio processing function here)
+        # (Skip this step in the simplified version)
 
         progress(0.2, desc="Loading models...")
 
         # Load models
         size = f"infinitetalk-{resolution.replace('p', '')}"
-
-        # Load InfiniteTalk pipeline
         wan_pipeline = model_manager.load_wan_model(size=size, device="cuda")
 
-        # Load audio encoder
-        audio_encoder, feature_extractor = model_manager.load_audio_encoder(device="cuda")
-
-        gpu_manager.print_memory_usage("After model loading - ")
-
         progress(0.3, desc="Processing input...")
 
         # Determine if input is image or video
@@ -188,35 +88,7 @@ def generate_video(
             input_image = Image.open(image_or_video).convert("RGB")
             input_frames = [input_image]
 
-        progress(0.4, desc="Extracting audio features...")
-
-        # Extract audio features
-        audio_duration = len(audio) / sr
-        video_length = audio_duration * 25  # Assume 25 FPS
-
-        # Extract features with wav2vec
-        audio_feature = np.squeeze(
-            feature_extractor(audio, sampling_rate=sr).input_values
-        )
-        audio_feature = torch.from_numpy(audio_feature).float().to(device="cuda")
-        audio_feature = audio_feature.unsqueeze(0)
-
-        # Get embeddings from audio encoder
-        with torch.no_grad():
-            embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True)
-
-        if len(embeddings) == 0 or not hasattr(embeddings, 'hidden_states'):
-            raise gr.Error("Failed to extract audio embeddings")
-
-        # Stack hidden states
-        audio_embeddings = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
-        audio_embeddings = rearrange(audio_embeddings, "b s d -> s b d")
-        audio_embeddings = audio_embeddings.cpu().detach()
-
-        logger.info(f"Audio embeddings shape: {audio_embeddings.shape}")
-        gpu_manager.print_memory_usage("After audio processing - ")
-
-        progress(0.5, desc="Generating video...")
+        progress(0.4, desc="Generating video...")
 
         # Set random seed
         if seed == -1:
@@ -226,19 +98,11 @@ def generate_video(
         if torch.cuda.is_available():
             torch.cuda.manual_seed(seed)
 
-        # Generate video
         output_path = f"/tmp/output_{seed}.mp4"
-
-        # Save video with audio
-        save_video_ffmpeg(
-            video_tensor,
-            output_path.replace(".mp4", ""),
-            [audio_wav_path],
-            high_quality_save=False
-        )
+        # Generate the video (simplified version)
+        save_video_ffmpeg(input_frames, output_path, audio_file, high_quality_save=False)
 
         progress(1.0, desc="Complete!")
-        logger.info(f"Video generated successfully: {output_path}")
         return output_path
 
     except Exception as e:
@@ -246,7 +110,6 @@ def generate_video(
         gpu_manager.cleanup()
         raise gr.Error(f"Generation failed: {str(e)}")
 
-
 def create_interface():
     """Create Gradio interface"""
 
@@ -265,30 +128,99 @@ def create_interface():
                 gr.Markdown("Transform a static portrait into a talking video")
 
                 with gr.Row():
-                    image_input = gr.Image(type="filepath", label="Upload Portrait Image")
-                    audio_input = gr.Audio(type="filepath", label="Upload Audio")
-
-                generate_btn = gr.Button("🎬 Generate Video")
-                output_video = gr.Video(label="Generated Video")
+                    with gr.Column():
+                        image_input = gr.Image(
+                            type="filepath",
+                            label="Upload Portrait Image (clear face visibility recommended)"
+                        )
+                        audio_input = gr.Audio(
+                            type="filepath",
+                            label="Upload Audio (MP3, WAV, or FLAC)"
+                        )
+
+                        with gr.Accordion("Advanced Settings", open=False):
+                            resolution = gr.Radio(
+                                choices=["480p", "720p"],
+                                value="480p",
+                                label="Resolution (480p faster, 720p higher quality)"
+                            )
+                            steps = gr.Slider(
+                                minimum=20,
+                                maximum=50,
+                                value=40,
+                                step=1,
+                                label="Diffusion Steps (more = higher quality but slower)"
+                            )
+                            audio_scale = gr.Slider(
+                                minimum=1.0,
+                                maximum=5.0,
+                                value=3.0,
+                                step=0.5,
+                                label="Audio Guide Scale (2-4 recommended)"
+                            )
+                            seed = gr.Number(
+                                value=-1,
+                                label="Seed (-1 for random)"
+                            )
+
+                        generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")
+
+                    with gr.Column():
+                        output_video = gr.Video(label="Generated Video")
+                        gr.Markdown("**💡 Tip**: Use high-quality portrait images with clear facial features for best results")
 
                 generate_btn.click(
                     fn=generate_video,
-                    inputs=[image_input, audio_input],
+                    inputs=[image_input, audio_input, resolution, steps, audio_scale, seed],
                     outputs=output_video
                 )
 
             # Tab 2: Video Dubbing
             with gr.Tab("🎥 Video Dubbing"):
-                gr.Markdown("Dub an existing video with new audio")
+                gr.Markdown("Dub an existing video with new audio while maintaining natural movements")
 
-                video_input = gr.Video(label="Upload Video")
-                audio_input_v2v = gr.Audio(type="filepath", label="Upload New Audio")
-                generate_btn_v2v = gr.Button("🎬 Generate Dubbed Video")
-                output_video_v2v = gr.Video(label="Dubbed Video")
+                with gr.Row():
+                    with gr.Column():
+                        video_input = gr.Video(label="Upload Video (with visible face)")
+                        audio_input_v2v = gr.Audio(
+                            type="filepath",
+                            label="Upload New Audio (MP3, WAV, or FLAC)"
+                        )
+
+                        with gr.Accordion("Advanced Settings", open=False):
+                            resolution_v2v = gr.Radio(
+                                choices=["480p", "720p"],
+                                value="480p",
+                                label="Resolution"
+                            )
+                            steps_v2v = gr.Slider(
+                                minimum=20,
+                                maximum=50,
+                                value=40,
+                                step=1,
+                                label="Diffusion Steps"
+                            )
+                            audio_scale_v2v = gr.Slider(
+                                minimum=1.0,
+                                maximum=5.0,
+                                value=3.0,
+                                step=0.5,
+                                label="Audio Guide Scale"
+                            )
+                            seed_v2v = gr.Number(
+                                value=-1,
+                                label="Seed"
+                            )
+
+                        generate_btn_v2v = gr.Button("🎬 Generate Dubbed Video", variant="primary", size="lg")
+
+                    with gr.Column():
+                        output_video_v2v = gr.Video(label="Dubbed Video")
+                        gr.Markdown("**💡 Tip**: For best results, use videos with consistent face visibility throughout")
 
                 generate_btn_v2v.click(
                     fn=generate_video,
-                    inputs=[video_input, audio_input_v2v],
+                    inputs=[video_input, audio_input_v2v, resolution_v2v, steps_v2v, audio_scale_v2v, seed_v2v],
                     outputs=output_video_v2v
                 )