app.py

import gradio as gr
import torch
import os
import gc
import numpy as np
import tempfile
from typing import Optional, Tuple
import time

# ZeroGPU with H200 support
try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ Spaces library loaded - H200 detected!")
except ImportError:
    SPACES_AVAILABLE = False
    class spaces:
        @staticmethod
        def GPU(duration=300):
            def decorator(func): return func
            return decorator

# Environment
IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") == "true"
IS_SPACES = os.environ.get("SPACE_ID") is not None
HAS_CUDA = torch.cuda.is_available()

print(f"🚀 H200 Environment: ZeroGPU={IS_ZERO_GPU}, Spaces={IS_SPACES}, CUDA={HAS_CUDA}")

# Premium models optimized for H200's massive memory
PREMIUM_MODELS = {
    "ltx": {
        "id": "Lightricks/LTX-Video",
        "name": "LTX-Video",
        "pipeline_class": "LTXVideoPipeline",
        "resolution_options": [(512, 512), (768, 768), (1024, 1024), (1280, 720), (1920, 1080)],
        "max_frames": 161,  # H200 can handle more frames
        "dtype": torch.bfloat16,
        "priority": 1,
        "description": "Lightricks' flagship model - professional quality"
    },
    "hunyuan": {
        "id": "tencent/HunyuanVideo", 
        "name": "HunyuanVideo",
        "pipeline_class": "HunyuanVideoPipeline",
        "resolution_options": [(512, 512), (768, 768), (1024, 1024), (1280, 720)],
        "max_frames": 129,  # Extended for H200
        "dtype": torch.bfloat16,
        "priority": 2,
        "description": "Tencent's advanced video model with superior motion"
    },
    "wan": {
        "id": "wangfuyun/AnimateLCM",
        "name": "AnimateLCM", 
        "pipeline_class": "DiffusionPipeline",
        "resolution_options": [(512, 512), (768, 768), (1024, 1024)],
        "max_frames": 64,
        "dtype": torch.float16,
        "priority": 3,
        "description": "Fast, high-quality animation model"
    },
    "cogvideo": {
        "id": "THUDM/CogVideoX-5b",
        "name": "CogVideoX-5B",
        "pipeline_class": "CogVideoXPipeline", 
        "resolution_options": [(720, 480), (1280, 720)],
        "max_frames": 49,
        "dtype": torch.bfloat16,
        "priority": 4,
        "description": "CogVideo's 5B parameter model"
    }
}

# Global variables
MODEL = None
MODEL_INFO = None
LOADING_ERROR = None

def get_gpu_memory():
    """Get H200 GPU memory info"""
    if HAS_CUDA:
        try:
            total_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            allocated = torch.cuda.memory_allocated(0) / (1024**3)
            cached = torch.cuda.memory_reserved(0) / (1024**3)
            return total_memory, allocated, cached
        except:
            return 0, 0, 0
    return 0, 0, 0

def load_premium_model():
    """Load first available premium model with H200 optimizations"""
    global MODEL, MODEL_INFO, LOADING_ERROR
    
    if MODEL is not None:
        return True
    
    # Sort models by priority
    sorted_models = sorted(PREMIUM_MODELS.items(), key=lambda x: x[1]["priority"])
    
    for key, info in sorted_models:
        try:
            print(f"🔄 Loading {info['name']} on H200...")
            total_mem, allocated, cached = get_gpu_memory()
            print(f"💾 GPU Memory: {total_mem:.1f}GB total, {allocated:.1f}GB allocated")
            
            from diffusers import DiffusionPipeline
            
            # Try specific pipeline class first
            try:
                if info["pipeline_class"] == "LTXVideoPipeline":
                    from diffusers import LTXVideoPipeline
                    pipe = LTXVideoPipeline.from_pretrained(
                        info["id"],
                        torch_dtype=info["dtype"],
                        use_safetensors=True,
                        variant="fp16"
                    )
                elif info["pipeline_class"] == "HunyuanVideoPipeline":
                    from diffusers import HunyuanVideoPipeline  
                    pipe = HunyuanVideoPipeline.from_pretrained(
                        info["id"],
                        torch_dtype=info["dtype"],
                        use_safetensors=True,
                        variant="fp16"
                    )
                elif info["pipeline_class"] == "CogVideoXPipeline":
                    from diffusers import CogVideoXPipeline
                    pipe = CogVideoXPipeline.from_pretrained(
                        info["id"],
                        torch_dtype=info["dtype"],
                        use_safetensors=True
                    )
                else:
                    # Generic DiffusionPipeline
                    pipe = DiffusionPipeline.from_pretrained(
                        info["id"],
                        torch_dtype=info["dtype"],
                        use_safetensors=True,
                        variant="fp16",
                        trust_remote_code=True
                    )
            except ImportError as e:
                print(f"⚠️ Specific pipeline not available: {e}")
                print("Trying generic DiffusionPipeline...")
                pipe = DiffusionPipeline.from_pretrained(
                    info["id"],
                    torch_dtype=info["dtype"],
                    use_safetensors=True,
                    variant="fp16",
                    trust_remote_code=True
                )
            
            # H200 optimizations - we have plenty of memory!
            if HAS_CUDA:
                pipe = pipe.to("cuda")
                print(f"📱 Moved {info['name']} to H200 CUDA")
                
                # Enable all optimizations but keep model in VRAM
                if hasattr(pipe, 'enable_vae_slicing'):
                    pipe.enable_vae_slicing()
                if hasattr(pipe, 'enable_vae_tiling'):
                    pipe.enable_vae_tiling()
                if hasattr(pipe, 'enable_memory_efficient_attention'):
                    pipe.enable_memory_efficient_attention()
                # Don't use CPU offload on H200 - keep everything in GPU memory
                
                # Enable xformers if available for extra speed
                try:
                    pipe.enable_xformers_memory_efficient_attention()
                    print("🚀 XFormers acceleration enabled")
                except:
                    print("⚠️ XFormers not available")
            
            MODEL = pipe
            MODEL_INFO = info
            
            final_mem = torch.cuda.memory_allocated(0) / (1024**3)
            print(f"✅ {info['name']} loaded successfully! Memory used: {final_mem:.1f}GB")
            return True
            
        except Exception as e:
            print(f"❌ Failed to load {info['name']}: {e}")
            # Clear memory before trying next model
            if HAS_CUDA:
                torch.cuda.empty_cache()
            gc.collect()
            continue
    
    LOADING_ERROR = "All premium models failed to load"
    return False

@spaces.GPU(duration=300) if SPACES_AVAILABLE else lambda x: x  # 5 minutes for H200
def generate_video(
    prompt: str,
    negative_prompt: str = "",
    num_frames: int = 49,
    resolution: str = "1024x1024",
    num_inference_steps: int = 30,
    guidance_scale: float = 7.5,
    seed: int = -1,
    fps: int = 8
) -> Tuple[Optional[str], str]:
    """Generate premium video with H200 power"""
    
    global MODEL, MODEL_INFO, LOADING_ERROR
    
    # Load model if needed
    if not load_premium_model():
        return None, f"❌ No premium models available: {LOADING_ERROR}"
    
    # Input validation
    if not prompt.strip():
        return None, "❌ Please enter a valid prompt."
    
    if len(prompt) > 1000:  # H200 can handle longer prompts
        return None, "❌ Prompt too long. Please keep it under 1000 characters."
    
    # Parse resolution
    try:
        width, height = map(int, resolution.split('x'))
    except:
        width, height = 1024, 1024
    
    # Validate parameters against model capabilities  
    max_frames = MODEL_INFO["max_frames"]
    num_frames = min(max(num_frames, 8), max_frames)
    
    # Check if resolution is supported
    if (width, height) not in MODEL_INFO["resolution_options"]:
        # Use best supported resolution
        best_res = MODEL_INFO["resolution_options"][-1]  # Highest resolution
        width, height = best_res
        print(f"⚠️ Adjusted resolution to {width}x{height}")
    
    try:
        # H200 memory management - we have tons of memory!
        start_memory = torch.cuda.memory_allocated(0) / (1024**3) if HAS_CUDA else 0
        
        # Set seed
        if seed == -1:
            seed = np.random.randint(0, 2**32 - 1)
        
        device = "cuda" if HAS_CUDA else "cpu"
        generator = torch.Generator(device=device).manual_seed(seed)
        
        print(f"🎬 H200 Generation: {MODEL_INFO['name']} - '{prompt[:70]}...'")
        print(f"📐 {width}x{height}, {num_frames} frames, {num_inference_steps} steps")
        start_time = time.time()
        
        # Generate with H200's full power
        with torch.autocast(device, dtype=MODEL_INFO["dtype"]):
            generation_kwargs = {
                "prompt": prompt,
                "num_frames": num_frames,
                "height": height,
                "width": width,
                "num_inference_steps": num_inference_steps,
                "guidance_scale": guidance_scale,
                "generator": generator
            }
            
            # Add negative prompt if provided
            if negative_prompt.strip():
                generation_kwargs["negative_prompt"] = negative_prompt
            
            # Model-specific parameters
            if MODEL_INFO["name"] == "CogVideoX-5B":
                generation_kwargs["num_videos_per_prompt"] = 1
            
            # Generate with progress tracking
            print("🚀 Starting generation on H200...")
            result = MODEL(**generation_kwargs)
        
        end_time = time.time()
        generation_time = end_time - start_time
        
        # Extract video frames
        if hasattr(result, 'frames'):
            video_frames = result.frames[0]
        elif hasattr(result, 'videos'):
            video_frames = result.videos[0]
        else:
            return None, "❌ Could not extract video frames from model output"
        
        # Export with custom FPS
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_file:
            from diffusers.utils import export_to_video
            export_to_video(video_frames, tmp_file.name, fps=fps)
            video_path = tmp_file.name
        
        # Memory stats
        end_memory = torch.cuda.memory_allocated(0) / (1024**3) if HAS_CUDA else 0
        memory_used = end_memory - start_memory
        
        success_msg = f"""✅ **H200 Premium Video Generated!**

🚀 **Model:** {MODEL_INFO['name']} 
📝 **Prompt:** {prompt}
🎬 **Frames:** {num_frames} @ {fps} FPS
📐 **Resolution:** {width}x{height}
⚙️ **Inference Steps:** {num_inference_steps}
🎯 **Guidance Scale:** {guidance_scale}
🎲 **Seed:** {seed}
⏱️ **Generation Time:** {generation_time:.1f}s
🖥️ **Device:** H200 CUDA
💾 **Memory Used:** {memory_used:.1f}GB
🎥 **Video Length:** {num_frames/fps:.1f}s"""
        
        return video_path, success_msg
        
    except torch.cuda.OutOfMemoryError:
        # Should be rare on H200!
        torch.cuda.empty_cache()
        gc.collect()
        return None, "❌ GPU memory exceeded (rare on H200!). Try reducing parameters."
    
    except Exception as e:
        if HAS_CUDA:
            torch.cuda.empty_cache()
        gc.collect()
        return None, f"❌ Generation failed: {str(e)}"

def get_h200_status():
    """Get H200 specific status"""
    if not HAS_CUDA:
        return "❌ CUDA not available"
    
    try:
        total_mem, allocated, cached = get_gpu_memory()
        gpu_name = torch.cuda.get_device_name(0)
        
        model_status = "⏳ Model will load on first use"
        if MODEL is not None:
            model_status = f"✅ {MODEL_INFO['name']} loaded and ready"
        elif LOADING_ERROR:
            model_status = f"❌ {LOADING_ERROR}"
        
        return f"""## 🚀 H200 Status

**🖥️ Hardware:**
- GPU: {gpu_name}
- Total Memory: {total_mem:.1f} GB
- Allocated: {allocated:.1f} GB  
- Cached: {cached:.1f} GB
- Free: {total_mem - allocated:.1f} GB

**🤖 Model Status:**
{model_status}

**⚡ H200 Advantages:**
- 141GB HBM3 memory (3.5x more than A100)
- 4.8TB/s memory bandwidth
- Can handle larger models & longer videos
- Multiple high-res generations without swapping"""
        
    except Exception as e:
        return f"❌ Error getting H200 status: {e}"

def suggest_h200_settings():
    """Suggest optimal settings for H200"""
    if MODEL is None:
        return "Load a model first to get personalized recommendations"
    
    model_name = MODEL_INFO['name']
    max_frames = MODEL_INFO['max_frames']
    max_res = MODEL_INFO['resolution_options'][-1]
    
    return f"""## 🎯 H200 Optimized Settings for {model_name}

**🚀 Maximum Quality (Recommended):**
- Resolution: {max_res[0]}x{max_res[1]}
- Frames: {max_frames}
- Inference Steps: 50
- Expected Time: 3-5 minutes

**⚖️ Balanced (Fast & Good):**
- Resolution: 1024x1024
- Frames: {max_frames//2}
- Inference Steps: 30
- Expected Time: 1-2 minutes

**⚡ Speed Test:**
- Resolution: 512x512  
- Frames: 25
- Inference Steps: 20
- Expected Time: 30-60 seconds

**💡 H200 Tips:**
- Use longer prompts - H200 can handle complexity
- Try higher inference steps (30-50) for maximum quality
- Experiment with longer videos (40+ frames)
- Multiple generations won't cause memory issues"""

# Create H200-optimized interface
with gr.Blocks(title="H200 Premium Video Generator", theme=gr.themes.Glass()) as demo:
    
    gr.Markdown("""
    # 🚀 H200 Premium Video Generator
    
    **Powered by NVIDIA H200** • **141GB Memory** • **Premium Models Only**
    
    *LTX-Video • HunyuanVideo • CogVideoX-5B • AnimateLCM*
    """)
    
    # H200 status bar
    with gr.Row():
        gr.Markdown("""
        <div style="text-align: center; padding: 10px; background: linear-gradient(45deg, #FF6B6B, #4ECDC4); border-radius: 10px; color: white; font-weight: bold;">
        🔥 H200 ACTIVE - MAXIMUM PERFORMANCE MODE 🔥
        </div>
        """)
    
    with gr.Tab("🎥 H200 Video Generation"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="📝 Detailed Video Prompt (H200 can handle complexity!)",
                    placeholder="A breathtaking aerial view of a majestic golden eagle soaring through dramatic mountain peaks during a spectacular sunrise, with volumetric lighting piercing through morning mist, cinematic composition with dynamic camera movement following the eagle's graceful flight path, professional cinematography with shallow depth of field and warm golden color grading, 8K quality with film grain texture...",
                    lines=5,
                    max_lines=8
                )
                
                negative_prompt_input = gr.Textbox(
                    label="🚫 Negative Prompt",
                    placeholder="blurry, low quality, distorted, pixelated, watermark, text, signature, amateur, static, boring, unnatural motion...",
                    lines=2
                )
                
                with gr.Accordion("🚀 H200 Advanced Settings", open=True):
                    with gr.Row():
                        num_frames = gr.Slider(
                            minimum=8,
                            maximum=161,  # H200 can handle more
                            value=49,
                            step=1,
                            label="🎬 Frames (H200 can handle long videos!)"
                        )
                        
                        fps = gr.Slider(
                            minimum=4,
                            maximum=30,
                            value=8,
                            step=1,
                            label="🎞️ FPS (frames per second)"
                        )
                    
                    with gr.Row():
                        resolution = gr.Dropdown(
                            choices=["512x512", "768x768", "1024x1024", "1280x720", "1920x1080"],
                            value="1024x1024",
                            label="📐 Resolution (H200 loves high-res!)"
                        )
                        
                        num_steps = gr.Slider(
                            minimum=15,
                            maximum=100,  # H200 can handle more steps
                            value=30,
                            step=1,
                            label="⚙️ Inference Steps (more = better quality)"
                        )
                    
                    with gr.Row():
                        guidance_scale = gr.Slider(
                            minimum=1.0,
                            maximum=20.0,
                            value=7.5,
                            step=0.5,
                            label="🎯 Guidance Scale"
                        )
                        
                        seed = gr.Number(
                            label="🎲 Seed (-1 for random)",
                            value=-1,
                            precision=0
                        )
                
                generate_btn = gr.Button(
                    "🚀 Generate on H200", 
                    variant="primary", 
                    size="lg"
                )
                
                gr.Markdown("""
                **⏱️ H200 Generation:** 1-5 minutes depending on settings
                
                **🔥 H200 Power:**
                - 141GB memory = No limits!
                - Generate 1080p videos
                - 100+ frames possible  
                - 50+ inference steps for max quality
                """)
                
            with gr.Column(scale=1):
                video_output = gr.Video(
                    label="🎥 H200 Generated Premium Video",
                    height=400
                )
                
                result_text = gr.Textbox(
                    label="📋 H200 Generation Report",
                    lines=12,
                    show_copy_button=True
                )
        
        # Event handler
        generate_btn.click(
            fn=generate_video,
            inputs=[
                prompt_input, negative_prompt_input, num_frames,
                resolution, num_steps, guidance_scale, seed, fps
            ],
            outputs=[video_output, result_text]
        )
        
        # H200-optimized examples
        gr.Examples(
            examples=[
                [
                    "A majestic golden eagle soaring through misty mountain peaks at sunrise, cinematic aerial cinematography with dramatic volumetric lighting, professional color grading with warm golden tones, shallow depth of field, dynamic camera movement tracking the eagle's flight, 8K quality with film grain",
                    "blurry, low quality, pixelated, static, amateur, watermark, text",
                    49, "1024x1024", 35, 7.5, 42, 8
                ],
                [
                    "Powerful ocean waves crashing against dramatic coastal cliffs during a storm, slow motion macro cinematography capturing water droplets and spray, dynamic lighting with storm clouds, professional cinematography with high contrast and desaturated colors",
                    "calm, peaceful, low quality, distorted, pixelated, watermark",
                    65, "1280x720", 40, 8.0, 123, 12
                ],
                [
                    "A steaming artisanal coffee cup on rustic wooden table by rain-streaked window, cozy cafe atmosphere with warm ambient lighting, shallow depth of field bokeh background, steam rising elegantly, cinematic close-up with perfect exposure",
                    "cold, harsh lighting, plastic, fake, low quality, blurry, text",
                    33, "1024x1024", 30, 7.0, 456, 8
                ],
                [
                    "Cherry blossom petals falling like snow in traditional Japanese garden with wooden bridge over koi pond, peaceful zen atmosphere with soft natural lighting, time-lapse effect showing seasonal transition, cinematic wide shot with perfect composition",
                    "modern, urban, chaotic, low quality, distorted, artificial, watermark",
                    81, "1280x720", 45, 7.5, 789, 10
                ]
            ],
            inputs=[prompt_input, negative_prompt_input, num_frames, resolution, num_steps, guidance_scale, seed, fps]
        )
    
    with gr.Tab("💾 H200 Status"):
        with gr.Row():
            status_btn = gr.Button("🔍 Check H200 Status", variant="secondary")
            settings_btn = gr.Button("🎯 Get Optimal Settings", variant="secondary")
            
        status_output = gr.Markdown()
        settings_output = gr.Markdown()
        
        status_btn.click(fn=get_h200_status, outputs=status_output)
        settings_btn.click(fn=suggest_h200_settings, outputs=settings_output)
        
        # Auto-load status
        demo.load(fn=get_h200_status, outputs=status_output)
    
    with gr.Tab("🎬 H200 Master Guide"):
        gr.Markdown("""
        ## 🚀 H200 Video Generation Mastery
        
        ### 💎 Why H200 is Game-Changing:
        
        **🔥 Raw Power:**
        - **141GB HBM3 Memory** (vs 80GB A100)
        - **4.8TB/s Bandwidth** (vs 3.35TB/s A100)  
        - **67% More Memory** for bigger models & longer videos
        - **No Memory Swapping** = Consistent performance
        
        ### 🎯 H200-Optimized Strategies:
        
        **🎬 Long-Form Content (H200 Specialty):**
        - Frames: 80-161 (2-20 second videos)
        - Resolution: 1280x720 or 1024x1024
        - Steps: 40-50 for cinematic quality
        - Perfect for: Storytelling, commercials, art pieces
        
        **🖼️ Ultra High-Res (H200 Advantage):**
        - Resolution: 1920x1080 (if model supports)
        - Frames: 25-49 (manageable length)
        - Steps: 30-40
        - Perfect for: Wallpapers, presentations, demos
        
        **⚡ Rapid Prototyping:**
        - Multiple quick generations to test ideas
        - 512x512, 25 frames, 20 steps
        - Iterate quickly, then scale up
        
        ### ✍️ Advanced Prompt Engineering for H200:
        
        **Complex Scene Composition:**
        ```
        [Main Subject] + [Detailed Action] + [Environment Description] + 
        [Camera Work] + [Lighting] + [Color Grading] + [Technical Quality]
        ```
        
        **Example Structure:**
        - **Subject:** "A majestic red dragon"
        - **Action:** "gracefully flying through ancient mountain peaks"  
        - **Environment:** "shrouded in mystical fog with ancient ruins visible below"
        - **Camera:** "cinematic aerial tracking shot with dynamic movement"
        - **Lighting:** "golden hour lighting with volumetric rays piercing the mist"
        - **Grading:** "warm color palette with high contrast and film grain"
        - **Quality:** "8K cinematography with shallow depth of field"
        
        ### 🎨 Style Modifiers for Premium Results:
        
        **Cinematic Styles:**
        - "Christopher Nolan cinematography"
        - "Blade Runner 2049 aesthetic"  
        - "Studio Ghibli animation style"
        - "BBC Planet Earth documentary style"
        - "Marvel movie action sequence"
        
        **Technical Quality:**
        - "8K RED camera footage"
        - "IMAX quality cinematography" 
        - "Zeiss lens bokeh"
        - "Professional color grading"
        - "Film grain texture overlay"
        
        ### 🔧 H200 Performance Optimization:
        
        **Memory Management:**
        - H200's 141GB means you rarely hit limits
        - Can run multiple models simultaneously
        - No need for CPU offloading
        - Keep all components in GPU memory
        
        **Speed Optimization:**
        - Use bfloat16 for modern models (LTX, HunyuanVideo)
        - Enable XFormers attention for 20-30% speedup
        - Batch operations when possible
        - H200's bandwidth handles large tensors efficiently
        
        **Quality Maximization:**
        - Push inference steps to 40-50
        - Use guidance scales 7-12 for detailed control
        - Experiment with longer sequences (80+ frames)
        - Try ultra-high resolutions (1080p+)
        
        ### 🎪 Advanced Techniques:
        
        **Multi-Shot Sequences:**
        1. Generate wide establishing shot
        2. Generate medium character shot
        3. Generate close-up detail shot
        4. Combine in post-production
        
        **Style Consistency:**
        - Use same seed across generations
        - Maintain lighting keywords
        - Keep camera angle descriptions similar
        - Use consistent color palette terms
        
        **Temporal Coherence:**
        - Describe smooth motions
        - Avoid jump cuts in single prompts
        - Use transition words: "smoothly", "gradually", "continuously"
        - Specify motion speed: "slow motion", "time-lapse", "real-time"
        
        ### 🏆 H200 Best Practices:
        
        **DO:**
        ✅ Push the limits - H200 can handle complexity
        ✅ Use detailed, multi-sentence prompts
        ✅ Experiment with high frame counts
        ✅ Try maximum inference steps for quality
        ✅ Generate multiple variations quickly
        
        **DON'T:**
        ❌ Limit yourself to basic settings
        ❌ Worry about memory constraints  
        ❌ Skip negative prompts
        ❌ Use generic prompts
        ❌ Settle for low resolution
        
        ### 🎬 Genre-Specific Prompting:
        
        **Nature Documentary:**
        "BBC Planet Earth style, macro cinematography, natural lighting, wildlife behavior, David Attenborough quality"
        
        **Sci-Fi Epic:**
        "Blade Runner 2049 aesthetic, neon lighting, futuristic architecture, dramatic cinematography, cyberpunk atmosphere"
        
        **Fantasy Adventure:**  
        "Lord of the Rings cinematography, epic landscapes, mystical lighting, heroic composition, John Howe art style"
        
        **Commercial/Product:**
        "Apple commercial style, clean minimalist aesthetic, perfect lighting, premium quality, studio photography"
        
        Remember: H200's massive memory means you can be ambitious. Don't hold back! 🚀
        """)

# Launch with H200 optimizations
if __name__ == "__main__":
    demo.queue(max_size=3)  # Smaller queue for premium H200 generations
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        show_api=False
    )