media.py

# --- LIBRARIES ---
import torch
import gradio as gr
import random
import time
from diffusers import AutoPipelineForText2Image, TextToVideoSDPipeline
import gc
import os
import imageio

# --- AUTHENTICATION FOR HUGGING FACE SPACES ---
# This will read the token from a "Secret" you set in your Space's settings
# It's more secure and the correct way to do it on HF Spaces.
try:
    from huggingface_hub import login
    HF_TOKEN = os.environ.get('HF_TOKEN')
    if HF_TOKEN:
        login(token=HF_TOKEN)
        print("✅ Hugging Face Authentication successful.")
    else:
        print("⚠️ Hugging Face token not found in Space Secrets. Gated models may not be available.")
except ImportError:
    print("Could not import huggingface_hub. Please ensure it's in requirements.txt")

# --- CONFIGURATION & STATE ---
available_models = {
    "Fast Image (SDXL Turbo)": "stabilityai/sdxl-turbo",
    "Quality Image (SDXL)": "stabilityai/stable-diffusion-xl-base-1.0",
    "Video (Zeroscope)": "cerspense/zeroscope-v2-576w"
}
model_state = { "current_pipe": None, "loaded_model_name": None }


# --- CORE GENERATION FUNCTION ---
# This is a generator function, which yields updates to the UI.
def generate_media(model_key, prompt, negative_prompt, steps, cfg_scale, width, height, seed, num_frames):
    # --- Model Loading Logic ---
    # If the requested model isn't the one we have loaded, switch them.
    if model_state.get("loaded_model_name") != model_key:
        print(f"Switching to {model_key}. Unloading previous model...")
        yield {status_textbox: f"Unloading previous model..."} # UI Update
        if model_state.get("current_pipe"):
            del model_state["current_pipe"]
            gc.collect()
            torch.cuda.empty_cache()

        model_id = available_models[model_key]
        print(f"Loading {model_id}...")
        yield {status_textbox: f"Loading {model_id}... This can take a minute."} # UI Update

        # Load the correct pipeline based on model type
        if "Image" in model_key:
            pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16")
        elif "Video" in model_key:
            pipe = TextToVideoSDPipeline.from_pretrained(model_id, torch_dtype=torch.float16)

        pipe.to("cuda")
        # Offload larger models to save VRAM, but keep fast models fully on GPU
        if "Turbo" not in model_key and "Video" not in model_key:
            pipe.enable_model_cpu_offload()

        model_state["current_pipe"] = pipe
        model_state["loaded_model_name"] = model_key
        print("✅ Model loaded successfully.")

    pipe = model_state["current_pipe"]
    generator = torch.Generator("cuda").manual_seed(seed)
    yield {status_textbox: f"Generating with {model_key}..."} # UI Update

    # --- Generation Logic ---
    if "Image" in model_key:
        print("Generating image...")
        if "Turbo" in model_key: # Special settings for SDXL Turbo
            num_steps, guidance_scale = 1, 0.0
        else:
            num_steps, guidance_scale = int(steps), float(cfg_scale)

        image = pipe(
            prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_steps,
            guidance_scale=guidance_scale, width=int(width), height=int(height), generator=generator
        ).images[0]
        print("✅ Image generation complete.")
        yield {output_image: image, output_video: None, status_textbox: f"Seed used: {seed}"}

    elif "Video" in model_key:
        print("Generating video...")
        video_frames = pipe(prompt=prompt, num_inference_steps=int(steps), height=320, width=576, num_frames=int(num_frames), generator=generator).frames
        
        video_path = f"/tmp/video_{seed}.mp4"
        imageio.mimsave(video_path, video_frames, fps=12)
        print(f"✅ Video saved to {video_path}")
        yield {output_image: None, output_video: video_path, status_textbox: f"Seed used: {seed}"}


# --- GRADIO USER INTERFACE ---
with gr.Blocks(theme='gradio/soft') as demo:
    gr.Markdown("# The Generative Media Suite")
    gr.Markdown("Create fast images, high-quality images, or short videos. Created by cheeseman182.")
    seed_state = gr.State(-1)

    with gr.Row():
        with gr.Column(scale=2):
            model_selector = gr.Radio(label="Select Model", choices=list(available_models.keys()), value=list(available_models.keys())[0])
            prompt_input = gr.Textbox(label="Prompt", lines=4, placeholder="An astronaut riding a horse on Mars, cinematic...")
            negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=2, value="ugly, blurry, deformed, watermark, text")
            
            with gr.Accordion("Settings", open=True):
                steps_slider = gr.Slider(1, 100, 30, step=1, label="Inference Steps")
                cfg_slider = gr.Slider(0.0, 15.0, 7.5, step=0.5, label="Guidance Scale (CFG)")
                with gr.Row():
                    width_slider = gr.Slider(256, 1024, 768, step=64, label="Width")
                    height_slider = gr.Slider(256, 1024, 768, step=64, label="Height")
                num_frames_slider = gr.Slider(12, 48, 24, step=4, label="Video Frames", visible=False)
                seed_input = gr.Number(-1, label="Seed (-1 for random)")
            
            generate_button = gr.Button("Generate", variant="primary")

        with gr.Column(scale=3):
            output_image = gr.Image(label="Image Result", interactive=False, height="60vh", visible=True)
            output_video = gr.Video(label="Video Result", interactive=False, height="60vh", visible=False)
            status_textbox = gr.Textbox(label="Status", interactive=False)

    # --- UI Logic ---
    def update_ui_on_model_change(model_key):
        is_video = "Video" in model_key
        is_turbo = "Turbo" in model_key
        return {
            steps_slider: gr.update(interactive=not is_turbo, value=1 if is_turbo else 30),
            cfg_slider: gr.update(interactive=not is_turbo, value=0.0 if is_turbo else 7.5),
            width_slider: gr.update(visible=not is_video),
            height_slider: gr.update(visible=not is_video),
            num_frames_slider: gr.update(visible=is_video),
            output_image: gr.update(visible=not is_video),
            output_video: gr.update(visible=is_video)
        }
    model_selector.change(update_ui_on_model_change, model_selector, [steps_slider, cfg_slider, width_slider, height_slider, num_frames_slider, output_image, output_video])

    # --- Button Logic ---
    # This chain first sets the seed, then calls the main generation function.
    click_event = generate_button.click(
        fn=lambda s: (s if s != -1 else random.randint(0, 2**32 - 1)),
        inputs=seed_input,
        outputs=seed_state,
        queue=False
    ).then(
        fn=generate_media,
        inputs=[model_selector, prompt_input, negative_prompt_input, steps_slider, cfg_slider, width_slider, height_slider, seed_state, num_frames_slider],
        outputs=[output_image, output_video, status_textbox]
    )

# This is the correct way to launch on Hugging Face Spaces
demo.launch()