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 import
try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    class spaces:
        @staticmethod
        def GPU(duration=60):
            def decorator(func):
                return func
            return decorator

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

def load_ltx_model_manual():
    """Manually load LTX-Video model using transformers"""
    try:
        print("🔄 Attempting to load LTX-Video with transformers...")
        
        from transformers import AutoModel, AutoTokenizer, AutoProcessor
        
        model_id = "Lightricks/LTX-Video"
        
        # Try loading with AutoModel
        try:
            processor = AutoProcessor.from_pretrained(model_id)
            model = AutoModel.from_pretrained(
                model_id,
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
                trust_remote_code=True  # Important for new models
            )
            
            if torch.cuda.is_available():
                model = model.to("cuda")
            
            print("✅ Model loaded with transformers")
            return model, processor, None
            
        except Exception as e:
            print(f"AutoModel failed: {e}")
            return None, None, str(e)
            
    except Exception as e:
        return None, None, f"Manual loading failed: {e}"

def load_alternative_video_model():
    """Load a working alternative video generation model"""
    try:
        print("🔄 Loading alternative video model...")
        
        from diffusers import DiffusionPipeline
        
        # Use Zeroscope or ModelScope as alternatives
        alternatives = [
            "cerspense/zeroscope_v2_576w",
            "damo-vilab/text-to-video-ms-1.7b",
            "ali-vilab/text-to-video-ms-1.7b"
        ]
        
        for model_id in alternatives:
            try:
                print(f"Trying {model_id}...")
                pipe = DiffusionPipeline.from_pretrained(
                    model_id,
                    torch_dtype=torch.float16,
                    use_safetensors=True,
                    variant="fp16"
                )
                
                if torch.cuda.is_available():
                    pipe = pipe.to("cuda")
                
                # Enable optimizations
                pipe.enable_sequential_cpu_offload()
                pipe.enable_vae_slicing()
                
                print(f"✅ Successfully loaded {model_id}")
                return pipe, model_id, None
                
            except Exception as e:
                print(f"Failed to load {model_id}: {e}")
                continue
        
        return None, None, "All alternative models failed"
        
    except Exception as e:
        return None, None, f"Alternative loading failed: {e}"

def create_mock_video(prompt, num_frames=16, width=512, height=512):
    """Create a mock video for demonstration"""
    try:
        import cv2
        from PIL import Image, ImageDraw, ImageFont
        
        # Create temporary video file
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_file:
            video_path = tmp_file.name
        
        # Video settings
        fps = 8
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
        
        # Color themes
        colors = [(255, 100, 100), (100, 255, 100), (100, 100, 255), (255, 255, 100)]
        
        for i in range(num_frames):
            # Create frame
            img = Image.new('RGB', (width, height), color=colors[i % len(colors)])
            draw = ImageDraw.Draw(img)
            
            try:
                font = ImageFont.truetype("arial.ttf", 24)
            except:
                font = ImageFont.load_default()
            
            # Add text
            draw.text((50, height//2 - 50), f"Frame {i+1}/{num_frames}", fill=(255, 255, 255), font=font)
            draw.text((50, height//2), f"Prompt: {prompt[:30]}...", fill=(255, 255, 255), font=font)
            draw.text((50, height//2 + 50), "DEMO MODE", fill=(0, 0, 0), font=font)
            
            # Convert to OpenCV format
            frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            out.write(frame)
        
        out.release()
        return video_path
        
    except Exception as e:
        return None

# Global variables
MODEL = None
PROCESSOR = None
MODEL_TYPE = None
MODEL_ERROR = None

def initialize_model():
    """Initialize model with fallback options"""
    global MODEL, PROCESSOR, MODEL_TYPE, MODEL_ERROR
    
    if MODEL is not None:
        return True
    
    if MODEL_ERROR is not None:
        return False
    
    print("🚀 Initializing video model...")
    
    # Strategy 1: Try manual LTX-Video loading
    print("Trying LTX-Video...")
    MODEL, PROCESSOR, error = load_ltx_model_manual()
    if MODEL is not None:
        MODEL_TYPE = "LTX-Video"
        return True
    
    print(f"LTX-Video failed: {error}")
    
    # Strategy 2: Try alternative models
    print("Trying alternative models...")
    MODEL, MODEL_TYPE, error = load_alternative_video_model()
    if MODEL is not None:
        PROCESSOR = None  # Diffusion pipeline doesn't need separate processor
        return True
    
    print(f"Alternative models failed: {error}")
    
    # Strategy 3: Use mock generation
    MODEL_TYPE = "mock"
    MODEL_ERROR = "All models failed - using demo mode"
    return False

@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 with fallback strategies"""
    
    # Initialize model
    model_loaded = initialize_model()
    
    # Input validation
    if not prompt.strip():
        return None, "❌ Please enter a valid prompt."
    
    # Limit parameters
    num_frames = min(max(num_frames, 8), 25)
    num_inference_steps = min(max(num_inference_steps, 10), 30)
    height = min(max(height, 256), 768)
    width = min(max(width, 256), 768)
    
    # Set seed
    if seed == -1:
        seed = np.random.randint(0, 2**32 - 1)
    
    try:
        # Clear memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
        
        start_time = time.time()
        
        if MODEL_TYPE == "mock" or not model_loaded:
            # Mock generation
            print("🎭 Using mock generation")
            video_path = create_mock_video(prompt, num_frames, width, height)
            
            if video_path:
                end_time = time.time()
                return video_path, f"""
🎭 **Demo Video Generated**

📝 Prompt: {prompt}
⚠️ Note: This is a demo mode because video models couldn't be loaded.

🎬 Frames: {num_frames}
📐 Resolution: {width}x{height}
⏱️ Time: {end_time - start_time:.1f}s
🔧 Status: {MODEL_ERROR or 'Demo mode'}

💡 **To enable real video generation:**
- Check if LTX-Video is available in your region
- Try upgrading diffusers: `pip install diffusers --upgrade`
- Or wait for official LTX-Video support in diffusers
                """
            else:
                return None, "❌ Even demo generation failed"
        
        elif MODEL_TYPE == "LTX-Video":
            # Manual LTX-Video generation
            print("🚀 Using LTX-Video")
            
            # This would need the actual implementation based on the model's API
            # For now, return a message about manual implementation needed
            return None, f"""
⚠️ **Manual Implementation Required**

LTX-Video model was loaded but requires custom generation code.
The model API is not yet standardized in diffusers.

📋 **Next Steps:**
1. Check Lightricks/LTX-Video model documentation
2. Implement custom inference pipeline
3. Or wait for official diffusers support

🔧 **Current Status:** Model loaded, awaiting implementation
            """
        
        else:
            # Alternative model generation
            print(f"🔄 Using {MODEL_TYPE}")
            
            generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
            
            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
            )
            
            # Export video
            video_frames = result.frames[0]
            
            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=8)
                video_path = tmp_file.name
            
            end_time = time.time()
            
            return video_path, f"""
✅ **Video Generated Successfully!**

📝 Prompt: {prompt}
🤖 Model: {MODEL_TYPE}
🎬 Frames: {num_frames}
📐 Resolution: {width}x{height}
⚙️ Steps: {num_inference_steps}
🎯 Guidance: {guidance_scale}
🎲 Seed: {seed}
⏱️ Time: {end_time - start_time:.1f}s
🖥️ Device: {'CUDA' if torch.cuda.is_available() else 'CPU'}
            """
    
    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 system information"""
    
    # Check what's available
    try:
        from diffusers import __version__ as diffusers_version
        available_pipelines = []
        try:
            from diffusers import LTXVideoPipeline
            available_pipelines.append("✅ LTXVideoPipeline")
        except ImportError:
            available_pipelines.append("❌ LTXVideoPipeline")
        
        try:
            from diffusers import DiffusionPipeline
            available_pipelines.append("✅ DiffusionPipeline")
        except ImportError:
            available_pipelines.append("❌ DiffusionPipeline")
            
    except ImportError:
        diffusers_version = "❌ Not installed"
        available_pipelines = ["❌ Diffusers not available"]
    
    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 '❌'}

**Packages:**
- PyTorch: {torch.__version__}
- Diffusers: {diffusers_version}
- Available Pipelines: {', '.join(available_pipelines)}

**Model Status:**
- Current Model: {MODEL_TYPE or 'Not loaded'}
- Status: {'✅ Ready' if MODEL is not None else '⚠️ ' + (MODEL_ERROR or 'Not initialized')}

**Recommendation:**
- LTX-Video is very new and may not be in stable diffusers yet
- Using alternative models or demo mode
- Check back later for official support
    """

# Create Gradio interface
with gr.Blocks(title="Video Generator with Fallbacks", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎬 Advanced Video Generator
    
    Attempts to use LTX-Video, falls back to alternative models, or provides demo mode.
    """)
    
    with gr.Tab("🎥 Generate Video"):
        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="📝 Video Prompt",
                    placeholder="A serene mountain lake at sunrise...",
                    lines=3
                )
                
                negative_prompt_input = gr.Textbox(
                    label="🚫 Negative Prompt",
                    placeholder="blurry, low quality...",
                    lines=2
                )
                
                with gr.Row():
                    num_frames = gr.Slider(8, 25, value=16, step=1, label="🎬 Frames")
                    num_steps = gr.Slider(10, 30, 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, 15.0, value=7.5, step=0.5, label="🎯 Guidance")
                    seed = gr.Number(value=-1, precision=0, label="🎲 Seed")
                
                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=8, show_copy_button=True)
        
        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]
        )
        
        gr.Examples(
            examples=[
                ["A peaceful cat in a sunny garden", "", 16, 512, 512, 20, 7.5, 42],
                ["Ocean waves at golden hour", "blurry", 20, 512, 512, 20, 8.0, 123],
                ["A butterfly on a flower", "", 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")
        system_output = gr.Markdown()
        
        info_btn.click(fn=get_system_info, outputs=system_output)
        demo.load(fn=get_system_info, outputs=system_output)

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