Upload app.py

app.py

@@ -0,0 +1,322 @@
+import gradio as gr
+import subprocess
+import os
+import tempfile
+import shutil
+from pathlib import Path
+import torch
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Constants
+DEFAULT_CONFIG_PATH = "configs/inference.yaml"
+DEFAULT_INPUT_FILE = "examples/infer_samples.txt"
+OUTPUT_DIR = Path("demo_out/gradio_outputs")
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+def generate_avatar_video(
+    reference_image,
+    audio_file, 
+    text_prompt,
+    seed=42,
+    num_steps=50,
+    guidance_scale=4.5,
+    audio_scale=None,
+    overlap_frames=13,
+    fps=25,
+    silence_duration=0.3,
+    resolution="720p",
+    progress=gr.Progress()
+):
+    """Generate an avatar video using OmniAvatar
+    
+    Args:
+        reference_image: Path to reference avatar image
+        audio_file: Path to audio file for lip sync
+        text_prompt: Text description of the video to generate
+        seed: Random seed for generation
+        num_steps: Number of inference steps
+        guidance_scale: Classifier-free guidance scale
+        audio_scale: Audio guidance scale (uses guidance_scale if None)
+        overlap_frames: Number of overlapping frames between chunks
+        fps: Frames per second
+        silence_duration: Duration of silence to add before/after audio
+        resolution: Output resolution ("480p" or "720p")
+        progress: Gradio progress callback
+        
+    Returns:
+        str: Path to generated video file
+    """
+    
+    try:
+        progress(0.1, desc="Preparing inputs")
+        
+        # Create temporary directory for this generation
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            
+            # Copy input files to temp directory
+            temp_image = temp_path / "input_image.jpeg"
+            temp_audio = temp_path / "input_audio.mp3"
+            shutil.copy(reference_image, temp_image)
+            shutil.copy(audio_file, temp_audio)
+            
+            # Create input file for inference script
+            input_file = temp_path / "input.txt"
+            # Format: prompt@@image_path@@audio_path
+            with open(input_file, 'w') as f:
+                f.write(f"{text_prompt}@@{temp_image}@@{temp_audio}\n")
+            
+            progress(0.2, desc="Configuring generation parameters")
+            
+            # Determine max_hw based on resolution
+            max_hw = 720 if resolution == "480p" else 1280
+            
+            # Build command to run inference script
+            cmd = [
+                "torchrun",
+                "--nproc_per_node=1",
+                "scripts/inference.py",
+                "--config", DEFAULT_CONFIG_PATH,
+                "--input_file", str(input_file),
+                "-hp", f"seed={seed},num_steps={num_steps},guidance_scale={guidance_scale},"
+                      f"overlap_frame={overlap_frames},fps={fps},silence_duration_s={silence_duration},"
+                      f"max_hw={max_hw},use_audio=True,i2v=True"
+            ]
+            
+            # Add audio scale if specified
+            if audio_scale is not None:
+                cmd[-1] += f",audio_scale={audio_scale}"
+            
+            progress(0.3, desc="Running OmniAvatar generation")
+            logger.info(f"Running command: {' '.join(cmd)}")
+            
+            # Run the inference script
+            env = os.environ.copy()
+            env['CUDA_VISIBLE_DEVICES'] = '0'  # Use first GPU
+            
+            process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                env=env
+            )
+            
+            # Monitor progress (simplified - in reality you'd parse the output)
+            stdout_lines = []
+            stderr_lines = []
+            
+            while True:
+                output = process.stdout.readline()
+                if output:
+                    stdout_lines.append(output.strip())
+                    logger.info(output.strip())
+                    
+                    # Update progress based on output
+                    if "Starting video generation" in output:
+                        progress(0.5, desc="Generating video frames")
+                    elif "[1/" in output:  # First chunk
+                        progress(0.6, desc="Processing video chunks")
+                    elif "Saving video" in output:
+                        progress(0.9, desc="Finalizing video")
+                
+                if process.poll() is not None:
+                    break
+            
+            # Get any remaining output
+            remaining_stdout, remaining_stderr = process.communicate()
+            if remaining_stdout:
+                stdout_lines.extend(remaining_stdout.strip().split('\n'))
+            if remaining_stderr:
+                stderr_lines.extend(remaining_stderr.strip().split('\n'))
+            
+            if process.returncode != 0:
+                error_msg = '\n'.join(stderr_lines)
+                logger.error(f"Inference failed with return code {process.returncode}")
+                logger.error(f"Error output: {error_msg}")
+                raise gr.Error(f"Video generation failed: {error_msg}")
+            
+            progress(0.95, desc="Retrieving generated video")
+            
+            # Find the generated video file
+            # The inference script saves to demo_out/{exp_name}/res_{input_file_name}_...
+            # We need to find the most recent video file
+            generated_videos = list(Path("demo_out").rglob("result_000.mp4"))
+            if not generated_videos:
+                raise gr.Error("No video file was generated")
+            
+            # Get the most recent video
+            latest_video = max(generated_videos, key=lambda p: p.stat().st_mtime)
+            
+            # Copy to output directory with unique name
+            output_filename = f"avatar_video_{os.getpid()}_{torch.randint(1000, 9999, (1,)).item()}.mp4"
+            output_path = OUTPUT_DIR / output_filename
+            shutil.copy(latest_video, output_path)
+            
+            progress(1.0, desc="Generation complete")
+            logger.info(f"Video saved to: {output_path}")
+            
+            return str(output_path)
+            
+    except Exception as e:
+        logger.error(f"Error generating video: {str(e)}")
+        raise gr.Error(f"Error generating video: {str(e)}")
+
+# Create the Gradio interface
+with gr.Blocks(title="OmniAvatar - Lipsynced Avatar Video Generation") as app:
+    gr.Markdown("""
+    # 🎭 OmniAvatar - Lipsynced Avatar Video Generation
+    
+    Generate videos with lipsynced avatars using a reference image and audio file.
+    Based on Wan2.1 with OmniAvatar enhancements for audio-driven avatar animation.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Input components
+            reference_image = gr.Image(
+                label="Reference Avatar Image",
+                type="filepath",
+                elem_id="reference_image"
+            )
+            
+            audio_file = gr.Audio(
+                label="Speech Audio File",
+                type="filepath",
+                elem_id="audio_file"
+            )
+            
+            text_prompt = gr.Textbox(
+                label="Video Description",
+                placeholder="Describe the video scene and actions...",
+                lines=3,
+                value="A person speaking naturally with subtle facial expressions"
+            )
+            
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=2147483647,
+                        step=1,
+                        value=42
+                    )
+                    
+                    resolution = gr.Radio(
+                        label="Resolution",
+                        choices=["480p", "720p"],
+                        value="720p"
+                    )
+                
+                with gr.Row():
+                    num_steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=10,
+                        maximum=100,
+                        step=5,
+                        value=50
+                    )
+                    
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.5,
+                        value=4.5
+                    )
+                
+                with gr.Row():
+                    audio_scale = gr.Slider(
+                        label="Audio Scale (leave 0 to use guidance scale)",
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.5,
+                        value=0.0
+                    )
+                    
+                    overlap_frames = gr.Slider(
+                        label="Overlap Frames",
+                        minimum=1,
+                        maximum=25,
+                        step=4,
+                        value=13,
+                        info="Must be 1 + 4*n"
+                    )
+                
+                with gr.Row():
+                    fps = gr.Slider(
+                        label="FPS",
+                        minimum=10,
+                        maximum=30,
+                        step=1,
+                        value=25
+                    )
+                    
+                    silence_duration = gr.Slider(
+                        label="Silence Duration (s)",
+                        minimum=0.0,
+                        maximum=2.0,
+                        step=0.1,
+                        value=0.3
+                    )
+            
+            generate_btn = gr.Button(
+                "🎬 Generate Avatar Video",
+                variant="primary"
+            )
+        
+        with gr.Column(scale=1):
+            # Output component
+            output_video = gr.Video(
+                label="Generated Avatar Video",
+                elem_id="output_video"
+            )
+            
+            # Examples
+            gr.Examples(
+                examples=[
+                    [
+                        "examples/images/0000.jpeg",
+                        "examples/audios/0000.MP3",
+                        "A professional woman giving a presentation with confident gestures"
+                    ],
+                ],
+                inputs=[reference_image, audio_file, text_prompt],
+                label="Example Inputs"
+            )
+    
+    # Connect the generate button
+    generate_btn.click(
+        fn=generate_avatar_video,
+        inputs=[
+            reference_image,
+            audio_file,
+            text_prompt,
+            seed,
+            num_steps,
+            guidance_scale,
+            audio_scale,
+            overlap_frames,
+            fps,
+            silence_duration,
+            resolution
+        ],
+        outputs=output_video
+    )
+    
+    gr.Markdown("""
+    ## 📝 Notes
+    - The reference image should be a clear frontal view of the person
+    - Audio should be clear speech without background music
+    - Generation may take several minutes depending on video length
+    - For best results, use high-quality input images and audio
+    """)
+
+# Launch the app
+if __name__ == "__main__":
+    app.launch(share=True)