Update app.py

app.py

@@ -5,12 +5,6 @@ Gradio Space for HuggingFace
 
 import os
 import sys
-
-# CRITICAL: Set environment variables BEFORE any torch/torchvision imports
-# This prevents torchvision from registering CUDA ops that don't exist at import time
-os.environ["TORCHVISION_DISABLE_META_REGISTRATIONS"] = "1"
-os.environ["TORCH_LOGS"] = "-all"  # Reduce torch logging noise
-
 import random
 import logging
 import warnings
@@ -19,43 +13,40 @@ from pathlib import Path
 import gradio as gr
 import torch
 import numpy as np
-
-# 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))
+import librosa
+import soundfile as sf
+import pyloudnorm as pyln
+from PIL import Image
+from einops import rearrange
 
 # Import utilities
 from utils.model_loader import ModelManager
 from utils.gpu_manager import gpu_manager
-
-# Import InfiniteTalk modules
 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
-
-# Audio processing
-import librosa
-import soundfile as sf
-import pyloudnorm as pyln
 from transformers import Wav2Vec2FeatureExtractor
 from src.audio_analysis.wav2vec2 import Wav2Vec2Model
 
-# Image/Video processing
-from PIL import Image
-from einops import rearrange
+# 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
 
-
 def initialize_models(progress=gr.Progress()):
     """Initialize models on first use"""
     global model_manager, models_loaded
@@ -98,16 +89,7 @@ def loudness_norm(audio_array, sr=16000, lufs=-20.0):
 
 
 def process_audio(audio_path, target_sr=16000):
-    """
-    Process audio file for InfiniteTalk (matches audio_prepare_single from reference)
-
-    Args:
-        audio_path: Path to audio file
-        target_sr: Target sample rate
-
-    Returns:
-        Processed audio array and sample rate
-    """
+    """Process audio file for InfiniteTalk"""
     try:
         # Load audio with librosa
         audio, sr = librosa.load(audio_path, sr=target_sr)
@@ -155,27 +137,12 @@ def generate_video(
     seed=-1,
     progress=gr.Progress()
 ):
-    """
-    Generate talking video from image or dub existing video
-
-    Args:
-        image_or_video: Input image or video file
-        audio_file: Audio file for lip-sync
-        resolution: Output resolution (480p or 720p)
-        steps: Number of diffusion steps
-        audio_guide_scale: Audio conditioning strength
-        seed: Random seed for reproducibility
-        progress: Gradio progress tracker
-
-    Returns:
-        Path to generated 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. "
-                "Please apply for a Community GPU Grant in the Space settings, or run this app locally with a GPU."
+                "⚠️ GPU not available. This Space requires GPU hardware to generate videos."
             )
 
         # Initialize models if needed
@@ -195,11 +162,6 @@ def generate_video(
         audio_duration = len(audio) / sr
         logger.info(f"Audio duration: {audio_duration:.2f}s")
 
-        # Calculate ZeroGPU duration
-        zerogpu_duration = gpu_manager.calculate_duration_for_zerogpu(
-            audio_duration, resolution
-        )
-
         progress(0.2, desc="Loading models...")
 
         # Load models
@@ -228,7 +190,7 @@ def generate_video(
 
         progress(0.4, desc="Extracting audio features...")
 
-        # Extract audio features (matches get_embedding from reference)
+        # Extract audio features
         audio_duration = len(audio) / sr
         video_length = audio_duration * 25  # Assume 25 FPS
 
@@ -246,8 +208,7 @@ def generate_video(
         if len(embeddings) == 0 or not hasattr(embeddings, 'hidden_states'):
             raise gr.Error("Failed to extract audio embeddings")
 
-        # Stack hidden states (matches reference implementation)
-        from einops import rearrange
+        # 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()
@@ -255,7 +216,7 @@ def generate_video(
         logger.info(f"Audio embeddings shape: {audio_embeddings.shape}")
         gpu_manager.print_memory_usage("After audio processing - ")
 
-        progress(0.5, desc="Generating video (this may take a minute)...")
+        progress(0.5, desc="Generating video...")
 
         # Set random seed
         if seed == -1:
@@ -265,69 +226,18 @@ def generate_video(
         if torch.cuda.is_available():
             torch.cuda.manual_seed(seed)
 
-        # Generate video with InfiniteTalk
+        # Generate video
         output_path = f"/tmp/output_{seed}.mp4"
 
-        # Prepare input for pipeline (following generate_infinitetalk.py structure)
-        with torch.no_grad():
-            logger.info(f"Generating {resolution} video with {steps} steps...")
-
-            # Save audio embeddings to temporary file (pipeline expects file path)
-            import tempfile
-            os.makedirs("/tmp/audio_embeddings", exist_ok=True)
-            emb_path = "/tmp/audio_embeddings/1.pt"
-            audio_wav_path = "/tmp/audio_embeddings/sum.wav"
-
-            torch.save(audio_embeddings, emb_path)
-            sf.write(audio_wav_path, audio, sr)
-
-            # Prepare input dictionary (matches generate_infinitetalk.py format)
-            input_clip = {
-                "prompt": "",  # Empty prompt for talking head
-                "cond_video": image_or_video,
-                "cond_audio": {
-                    "person1": emb_path
-                },
-                "video_audio": audio_wav_path
-            }
-
-            # Calculate sample_shift based on resolution
-            sample_shift = 7 if resolution == "480p" else 11
-
-            # Call InfiniteTalk pipeline
-            video_tensor = wan_pipeline.generate_infinitetalk(
-                input_clip,
-                size_buckget=size,
-                motion_frame=9,  # Default motion frame
-                frame_num=81,  # Default frame num (4n+1 format)
-                shift=sample_shift,
-                sampling_steps=steps,
-                text_guide_scale=5.0,  # Default text guidance
-                audio_guide_scale=audio_guide_scale,
-                seed=seed,
-                offload_model=True,
-                max_frames_num=81,  # For clip mode
-                color_correction_strength=1.0,
-                extra_args=None
-            )
-
-            # Save video with audio
-            from wan.utils.multitalk_utils import save_video_ffmpeg
-
-            save_video_ffmpeg(
-                video_tensor,
-                output_path.replace(".mp4", ""),  # Function adds .mp4 extension
-                [audio_wav_path],
-                high_quality_save=False
-            )
-
-        progress(0.9, desc="Finalizing...")
-
-        # Cleanup
-        gpu_manager.cleanup()
+        # Save video with audio
+        save_video_ffmpeg(
+            video_tensor,
+            output_path.replace(".mp4", ""),
+            [audio_wav_path],
+            high_quality_save=False
+        )
 
         progress(1.0, desc="Complete!")
-
         logger.info(f"Video generated successfully: {output_path}")
         return output_path
 
@@ -340,7 +250,7 @@ def generate_video(
 def create_interface():
     """Create Gradio interface"""
 
-    with gr.Blocks(title="InfiniteTalk - Talking Video Generator", theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(title="InfiniteTalk - Talking Video Generator") as demo:
         gr.Markdown("""
         # 🎬 InfiniteTalk - Talking Video Generator
 
@@ -355,101 +265,30 @@ def create_interface():
                 gr.Markdown("Transform a static portrait into a talking video")
 
                 with gr.Row():
-                    with gr.Column():
-                        image_input = gr.Image(
-                            type="filepath",
-                            label="Upload Portrait Image (clear face visibility recommended)"
-                        )
-                        audio_input_i2v = gr.Audio(
-                            type="filepath",
-                            label="Upload Audio (MP3, WAV, or FLAC)"
-                        )
-
-                        with gr.Accordion("Advanced Settings", open=False):
-                            resolution_i2v = gr.Radio(
-                                choices=["480p", "720p"],
-                                value="480p",
-                                label="Resolution (480p faster, 720p higher quality)"
-                            )
-                            steps_i2v = gr.Slider(
-                                minimum=20,
-                                maximum=50,
-                                value=40,
-                                step=1,
-                                label="Diffusion Steps (more = higher quality but slower)"
-                            )
-                            audio_scale_i2v = gr.Slider(
-                                minimum=1.0,
-                                maximum=5.0,
-                                value=3.0,
-                                step=0.5,
-                                label="Audio Guide Scale (2-4 recommended)"
-                            )
-                            seed_i2v = gr.Number(
-                                value=-1,
-                                label="Seed (-1 for random)"
-                            )
-
-                        generate_btn_i2v = gr.Button("🎬 Generate Video", variant="primary", size="lg")
-
-                    with gr.Column():
-                        output_video_i2v = gr.Video(label="Generated Video")
-                        gr.Markdown("**💡 Tip**: Use high-quality portrait images with clear facial features for best results")
-
-                generate_btn_i2v.click(
+                    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")
+
+                generate_btn.click(
                     fn=generate_video,
-                    inputs=[image_input, audio_input_i2v, resolution_i2v, steps_i2v, audio_scale_i2v, seed_i2v],
-                    outputs=output_video_i2v
+                    inputs=[image_input, audio_input],
+                    outputs=output_video
                 )
 
             # Tab 2: Video Dubbing
             with gr.Tab("🎥 Video Dubbing"):
-                gr.Markdown("Dub an existing video with new audio while maintaining natural movements")
+                gr.Markdown("Dub an existing video with new audio")
 
-                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")
+                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")
 
                 generate_btn_v2v.click(
                     fn=generate_video,
-                    inputs=[video_input, audio_input_v2v, resolution_v2v, steps_v2v, audio_scale_v2v, seed_v2v],
+                    inputs=[video_input, audio_input_v2v],
                     outputs=output_video_v2v
                 )
 
@@ -460,12 +299,6 @@ def create_interface():
         Powered by [InfiniteTalk](https://github.com/MeiGen-AI/InfiniteTalk) - Apache 2.0 License
 
         ⚠️ **Note**: This Space requires GPU hardware to generate videos. Apply for a Community GPU Grant in Settings.
-
-        💡 **Tips**:
-        - First generation downloads models (~15GB) and may take 2-3 minutes
-        - Use 480p for faster generation (~40s for 10s video)
-        - Use 720p for higher quality (slower but better results)
-        - Clear, well-lit images produce the best results
         """)
 
     return demo
@@ -473,5 +306,4 @@ def create_interface():
 
 if __name__ == "__main__":
     demo = create_interface()
-    demo.queue(max_size=10)
     demo.launch()