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
import subprocess
import sys

# ZeroGPU import
try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ Spaces library loaded successfully")
except ImportError:
    print("⚠️ Spaces library not available")
    SPACES_AVAILABLE = False
    # Create dummy decorator
    def spaces_gpu_decorator(duration=60):
        def decorator(func):
            return func
        return decorator
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()

# Environment checks
IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") == "true"
IS_SPACES = os.environ.get("SPACE_ID") is not None

print(f"Environment: ZeroGPU={IS_ZERO_GPU}, Spaces={IS_SPACES}")

def check_and_install_requirements():
    """Check and install missing requirements"""
    try:
        import diffusers
        print(f"✅ Diffusers version: {diffusers.__version__}")
        return True
    except ImportError:
        print("❌ Diffusers not found, attempting to install...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "diffusers[torch]>=0.30.0"])
            subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers>=4.35.0"])
            subprocess.check_call([sys.executable, "-m", "pip", "install", "accelerate"])
            import diffusers
            print(f"✅ Diffusers installed successfully: {diffusers.__version__}")
            return True
        except Exception as e:
            print(f"❌ Failed to install diffusers: {e}")
            return False

def load_model_safe():
    """Safely load the LTX-Video model with comprehensive error handling"""
    
    # First, ensure requirements are installed
    if not check_and_install_requirements():
        return None, "Failed to install required packages"
    
    try:
        print("🔄 Attempting to load LTX-Video model...")
        
        # Import after installation
        from diffusers import LTXVideoPipeline
        import torch
        
        model_id = "Lightricks/LTX-Video"
        
        # Check available memory
        if torch.cuda.is_available():
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            print(f"📊 Available GPU memory: {gpu_memory:.1f} GB")
        
        # Load with conservative settings
        print("📥 Loading pipeline...")
        pipe = LTXVideoPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            use_safetensors=True,
            variant="fp16"
        )
        
        # Move to GPU if available
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
            print("🚀 Model moved to GPU")
        
        # Enable optimizations
        try:
            pipe.enable_vae_slicing()
            pipe.enable_vae_tiling()
            print("⚡ Memory optimizations enabled")
        except Exception as e:
            print(f"⚠️ Some optimizations failed: {e}")
        
        print("✅ Model loaded successfully!")
        return pipe, None
        
    except ImportError as e:
        error_msg = f"Import error: {e}. Please check if diffusers is properly installed."
        print(f"❌ {error_msg}")
        return None, error_msg
        
    except Exception as e:
        error_msg = f"Model loading failed: {str(e)}"
        print(f"❌ {error_msg}")
        return None, error_msg

# Global model variable
MODEL = None
MODEL_ERROR = None

def initialize_model():
    """Initialize model on first use"""
    global MODEL, MODEL_ERROR
    if MODEL is None and MODEL_ERROR is None:
        print("🚀 Initializing model for first use...")
        MODEL, MODEL_ERROR = load_model_safe()
    return MODEL is not None

@spaces.GPU(duration=120) if SPACES_AVAILABLE else lambda x: x
def generate_video(
    prompt: str,
    negative_prompt: str = "",
    num_frames: int = 16,
    height: int = 512,
    width: int = 512,
    num_inference_steps: int = 20,
    guidance_scale: float = 7.5,
    seed: int = -1
) -> Tuple[Optional[str], str]:
    """Generate video using LTX-Video with ZeroGPU"""
    
    global MODEL, MODEL_ERROR
    
    # Initialize model if needed
    if not initialize_model():
        error_msg = f"❌ Model initialization failed: {MODEL_ERROR or 'Unknown error'}"
        return None, error_msg
    
    # Input validation
    if not prompt.strip():
        return None, "❌ Please enter a valid prompt."
    
    if len(prompt) > 200:
        return None, "❌ Prompt too long. Please keep it under 200 characters."
    
    # Limit parameters for stability
    num_frames = min(max(num_frames, 8), 24)
    num_inference_steps = min(max(num_inference_steps, 10), 25)
    height = min(max(height, 256), 768)
    width = min(max(width, 256), 768)
    
    try:
        # Clear memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        
        # Set seed
        if seed == -1:
            seed = np.random.randint(0, 2**32 - 1)
        
        generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
        
        print(f"🎬 Generating: '{prompt[:50]}...'")
        start_time = time.time()
        
        # Generate video
        with torch.autocast("cuda" if torch.cuda.is_available() else "cpu", dtype=torch.bfloat16):
            result = MODEL(
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt.strip() else None,
                num_frames=num_frames,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=generator,
            )
        
        end_time = time.time()
        generation_time = end_time - start_time
        
        # Save video
        video_frames = result.frames[0]
        
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_file:
            try:
                from diffusers.utils import export_to_video
                export_to_video(video_frames, tmp_file.name, fps=8)
                video_path = tmp_file.name
            except Exception as e:
                # Fallback: save as individual frames if export fails
                print(f"⚠️ Video export failed, trying alternative: {e}")
                return None, f"❌ Video export failed: {str(e)}"
        
        # Clear memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        
        success_msg = f"""✅ Video generated successfully!

📝 **Prompt:** {prompt}
🎬 **Frames:** {num_frames}
📐 **Resolution:** {width}x{height}
⚙️ **Inference Steps:** {num_inference_steps}
🎯 **Guidance Scale:** {guidance_scale}
🎲 **Seed:** {seed}
⏱️ **Generation Time:** {generation_time:.1f}s
🖥️ **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
⚡ **ZeroGPU:** {'✅' if IS_ZERO_GPU else '❌'}"""
        
        return video_path, success_msg
        
    except torch.cuda.OutOfMemoryError:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        return None, "❌ GPU memory exceeded. Try reducing frames/resolution or try again in a moment."
    
    except Exception as e:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        return None, f"❌ Generation failed: {str(e)}"

def get_system_info():
    """Get comprehensive system information"""
    
    # Check package versions
    package_info = {}
    try:
        import diffusers
        package_info['diffusers'] = diffusers.__version__
    except ImportError:
        package_info['diffusers'] = '❌ Not installed'
    
    try:
        import transformers
        package_info['transformers'] = transformers.__version__
    except ImportError:
        package_info['transformers'] = '❌ Not installed'
    
    # GPU info
    gpu_info = "❌ Not available"
    gpu_memory = 0
    if torch.cuda.is_available():
        try:
            gpu_info = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        except:
            gpu_info = "✅ Available (details unavailable)"
    
    return f"""## 🖥️ System Information

**Environment:**
- 🚀 ZeroGPU: {'✅ Active' if IS_ZERO_GPU else '❌ Not detected'}  
- 🏠 HF Spaces: {'✅' if IS_SPACES else '❌'}
- 🔥 CUDA: {'✅' if torch.cuda.is_available() else '❌'}
- 🖥️ GPU: {gpu_info} ({gpu_memory:.1f} GB)

**Packages:**
- PyTorch: {torch.__version__}
- Diffusers: {package_info.get('diffusers', 'Unknown')}
- Transformers: {package_info.get('transformers', 'Unknown')}
- Spaces: {'✅' if SPACES_AVAILABLE else '❌'}

**Model Status:**
- LTX-Video: {'✅ Loaded' if MODEL is not None else '⏳ Will load on first use' if MODEL_ERROR is None else f'❌ Error: {MODEL_ERROR}'}

**Tips:**
{'🎯 Ready to generate!' if MODEL is not None else '⚡ First generation will take longer due to model loading'}"""

def test_dependencies():
    """Test if all dependencies are working"""
    results = []
    
    # Test torch
    try:
        import torch
        results.append(f"✅ PyTorch {torch.__version__}")
        if torch.cuda.is_available():
            results.append(f"✅ CUDA {torch.version.cuda}")
        else:
            results.append("⚠️ CUDA not available")
    except Exception as e:
        results.append(f"❌ PyTorch: {e}")
    
    # Test diffusers
    try:
        import diffusers
        results.append(f"✅ Diffusers {diffusers.__version__}")
    except Exception as e:
        results.append(f"❌ Diffusers: {e}")
    
    # Test transformers
    try:
        import transformers
        results.append(f"✅ Transformers {transformers.__version__}")
    except Exception as e:
        results.append(f"❌ Transformers: {e}")
    
    return "\n".join(results)

# Create Gradio interface
with gr.Blocks(title="LTX-Video ZeroGPU", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🚀 LTX-Video Generator (ZeroGPU)
    
    Generate high-quality videos from text using **Lightricks LTX-Video** model with **ZeroGPU**!
    """)
    
    # Status indicator
    with gr.Row():
        gr.Markdown(f"""
        **Status:** {'🟢 ZeroGPU Active' if IS_ZERO_GPU else '🟡 CPU Mode'} | 
        **Environment:** {'HF Spaces' if IS_SPACES else 'Local'}
        """)
    
    with gr.Tab("🎥 Generate Video"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="📝 Video Prompt",
                    placeholder="A majestic eagle soaring through mountain peaks...",
                    lines=3,
                    max_lines=5
                )
                
                negative_prompt_input = gr.Textbox(
                    label="🚫 Negative Prompt (Optional)",
                    placeholder="blurry, low quality, distorted...",
                    lines=2
                )
                
                with gr.Accordion("⚙️ Settings", open=True):
                    with gr.Row():
                        num_frames = gr.Slider(8, 24, value=16, step=1, label="🎬 Frames")
                        num_steps = gr.Slider(10, 25, value=20, step=1, label="🔄 Steps") 
                    
                    with gr.Row():
                        width = gr.Dropdown([256, 512, 768], value=512, label="📐 Width")
                        height = gr.Dropdown([256, 512, 768], value=512, label="📏 Height")
                    
                    with gr.Row():
                        guidance_scale = gr.Slider(1.0, 12.0, value=7.5, step=0.5, label="🎯 Guidance")
                        seed = gr.Number(value=-1, precision=0, label="🎲 Seed (-1=random)")
                
                generate_btn = gr.Button("🚀 Generate Video", variant="primary", size="lg")
                
            with gr.Column(scale=1):
                video_output = gr.Video(label="🎥 Generated Video", height=400)
                result_text = gr.Textbox(label="📋 Results", lines=6, show_copy_button=True)
        
        # Event handlers
        generate_btn.click(
            fn=generate_video,
            inputs=[prompt_input, negative_prompt_input, num_frames, height, width, num_steps, guidance_scale, seed],
            outputs=[video_output, result_text]
        )
        
        # Examples
        gr.Examples(
            examples=[
                ["A peaceful cat sleeping in a sunny garden", "", 16, 512, 512, 20, 7.5, 42],
                ["Ocean waves at sunset, cinematic view", "blurry", 20, 512, 512, 20, 8.0, 123],
                ["A hummingbird hovering near red flowers", "", 16, 512, 512, 15, 7.0, 456]
            ],
            inputs=[prompt_input, negative_prompt_input, num_frames, height, width, num_steps, guidance_scale, seed]
        )
    
    with gr.Tab("ℹ️ System Info"):
        info_btn = gr.Button("🔍 Check System", variant="secondary")
        system_output = gr.Markdown()
        
        info_btn.click(fn=get_system_info, outputs=system_output)
        demo.load(fn=get_system_info, outputs=system_output)
    
    with gr.Tab("🔧 Debug"):
        test_btn = gr.Button("🧪 Test Dependencies")
        test_output = gr.Textbox(label="Test Results", lines=10)
        
        test_btn.click(fn=test_dependencies, outputs=test_output)

# Launch
if __name__ == "__main__":
    demo.queue(max_size=5)
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )