import gradio as gr
from transformers import pipeline

# Load the image and video generation pipelines
image_pipe = pipeline("text-to-image", model="stabilityai/sdxl", trust_remote_code=True)
video_pipe = pipeline("text-to-video", model="THUDM/CogVideoX1.5-5B-I2V", trust_remote_code=True)

def generate(prompt, media_type):
    if media_type == "image":
        image = image_pipe(prompt)[0]
        return image, None
    elif media_type == "video":
        video_path = video_pipe(prompt)[0]  # This returns a filepath
        return None, video_path

# Gradio interface
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe what you want to generate..."),
        gr.Radio(["image", "video"], label="Type")
    ],
    outputs=[
        gr.Image(type="pil", label="Generated Image"),
        gr.Video(label="Generated Video")
    ],
    title="🎨 NeuroMuse AI Generator",
    description="Generate stunning visuals or AI-powered video from your imagination using free models."
)

iface.launch()