import gradio as gr
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    # Initialize the diffusion pipeline
    logger.info("Loading diffusion pipeline...")
    pipe = DiffusionPipeline.from_pretrained(
        "heboya8/text2video-test-2",
        torch_dtype=torch.float16,
        trust_remote_code=True,
    )

    # Optimize for GPU memory
    logger.info("Enabling CPU offload and VAE slicing...")
    pipe.enable_model_cpu_offload()
    pipe.enable_vae_slicing()
except Exception as e:
    logger.error(f"Failed to initialize pipeline: {str(e)}")
    raise

def generate_video(prompt):
    try:
        logger.info(f"Generating video for prompt: {prompt}")
        # Generate video frames
        video_frames = pipe(
            prompt,
            num_inference_steps=50,  # Increased for better quality
            num_frames=16,          # Increased for longer video
        ).frames

        # Export frames to video file
        video_path = export_to_video(video_frames, output_video_path="output_video.mp4")
        logger.info(f"Video generated at: {video_path}")
        return video_path
    except Exception as e:
        logger.error(f"Error generating video: {str(e)}")
        return f"Error generating video: {str(e)}"

# Create Gradio interface
interface = gr.Interface(
    fn=generate_video,
    inputs=gr.Textbox(
        label="Enter your prompt",
        placeholder="e.g., a flower in a garden"
    ),
    outputs=gr.Video(label="Generated Video"),
    title="Text-to-Video Generator",
    description="Enter a text prompt to generate a video using the diffusion model."
)

# Launch the app
if __name__ == "__main__":
    logger.info("Launching Gradio interface...")
    interface.launch(server_name="0.0.0.0", server_port=7860)