app.py

import gradio as gr
import torch
import random
import hashlib
from diffusers import DiffusionPipeline
from transformers import pipeline
from diffusers.utils import export_to_video

# Optional: xformers optimization
try:
    import xformers
    has_xformers = True
except ImportError:
    has_xformers = False

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
MAX_SEED = 2**32 - 1

# Model lists ordered by size
image_models = {
    "Stable Diffusion 1.5 (light)": "runwayml/stable-diffusion-v1-5",
    "Stable Diffusion 2.1": "stabilityai/stable-diffusion-2-1",
    "Dreamlike 2.0": "dreamlike-art/dreamlike-photoreal-2.0",
    "Playground v2": "playgroundai/playground-v2-1024px-aesthetic",
    "Muse 512": "amused/muse-512-finetuned",
    "PixArt": "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
    "Kandinsky 3": "kandinsky-community/kandinsky-3",
    "BLIP Diffusion": "Salesforce/blipdiffusion",
    "SDXL Base 1.0 (heavy)": "stabilityai/stable-diffusion-xl-base-1.0",
    "OpenJourney (heavy)": "prompthero/openjourney"
}

text_models = {
    "GPT-2 (light)": "gpt2",
    "GPT-Neo 1.3B": "EleutherAI/gpt-neo-1.3B",
    "BLOOM 1.1B": "bigscience/bloom-1b1",
    "GPT-J 6B": "EleutherAI/gpt-j-6B",
    "Falcon 7B": "tiiuae/falcon-7b",
    "XGen 7B": "Salesforce/xgen-7b-8k-base",
    "BTLM 3B": "cerebras/btlm-3b-8k-base",
    "MPT 7B": "mosaicml/mpt-7b",
    "StableLM 2": "stabilityai/stablelm-2-1_6b",
    "LLaMA 2 7B (heavy)": "meta-llama/Llama-2-7b-hf"
}

video_models = {
    "CogVideoX-2B": "THUDM/CogVideoX-2b",
    "CogVideoX-5B": "THUDM/CogVideoX-5b",
    "AnimateDiff-Lightning": "ByteDance/AnimateDiff-Lightning",
    "ModelScope T2V": "damo-vilab/text-to-video-ms-1.7b",
    "VideoCrafter2": "VideoCrafter/VideoCrafter2",
    "Open-Sora-Plan-v1.2.0": "LanguageBind/Open-Sora-Plan-v1.2.0",
    "LTX-Video": "Lightricks/LTX-Video",
    "HunyuanVideo": "tencent/HunyuanVideo",
    "Latte-1": "maxin-cn/Latte-1",
    "LaVie": "Vchitect/LaVie"
}

# Caches
image_pipes = {}
text_pipes = {}
video_pipes = {}
image_cache = {}
text_cache = {}
video_cache = {}

def hash_inputs(*args):
    combined = "|".join(map(str, args))
    return hashlib.sha256(combined.encode()).hexdigest()

def generate_image(prompt, model_name, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    key = hash_inputs(prompt, model_name, seed)
    if key in image_cache:
        progress(100, desc="Using cached image.")
        return image_cache[key], seed

    progress(10, desc="Loading model...")
    if model_name not in image_pipes:
        pipe = DiffusionPipeline.from_pretrained(
            image_models[model_name],
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True
        )

        if torch.__version__.startswith("2"):
            pipe = torch.compile(pipe)
        if has_xformers and device == "cuda":
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except Exception:
                pass

        pipe.to(device)
        image_pipes[model_name] = pipe

    pipe = image_pipes[model_name]

    progress(40, desc="Generating image...")
    result = pipe(prompt=prompt, generator=torch.manual_seed(seed), num_inference_steps=15, width=512, height=512)
    image = result.images[0]
    image_cache[key] = image

    progress(100, desc="Done.")
    return image, seed

def generate_text(prompt, model_name, progress=gr.Progress(track_tqdm=True)):
    key = hash_inputs(prompt, model_name)
    if key in text_cache:
        progress(100, desc="Using cached text.")
        return text_cache[key]

    progress(10, desc="Loading model...")
    if model_name not in text_pipes:
        text_pipes[model_name] = pipeline(
            "text-generation",
            model=text_models[model_name],
            device=0 if device == "cuda" else -1
        )
    pipe = text_pipes[model_name]

    progress(40, desc="Generating text...")
    result = pipe(prompt, max_length=100, do_sample=True)[0]['generated_text']
    text_cache[key] = result

    progress(100, desc="Done.")
    return result

def generate_video(prompt, model_name, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    key = hash_inputs(prompt, model_name, seed)
    if key in video_cache:
        progress(100, desc="Using cached video.")
        return video_cache[key], seed

    progress(10, desc="Loading model...")
    if model_name not in video_pipes:
        pipe = DiffusionPipeline.from_pretrained(
            video_models[model_name],
            torch_dtype=torch_dtype,
            variant="fp16"
        )

        if torch.__version__.startswith("2"):
            pipe = torch.compile(pipe)
        if has_xformers and device == "cuda":
            try:
                pipe.enable_xformers_memory_efficient_attention()
            except Exception:
                pass

        pipe.to(device)
        video_pipes[model_name] = pipe

    pipe = video_pipes[model_name]

    progress(40, desc="Generating video...")
    result = pipe(prompt=prompt, generator=torch.manual_seed(seed), num_inference_steps=15)
    video_frames = result.frames[0]
    video_path = export_to_video(video_frames)
    video_cache[key] = video_path

    progress(100, desc="Done.")
    return video_path, seed

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# ⚡ Fast Multi-Model AI Playground with Caching")

    with gr.Tabs():
        # Image Generation
        with gr.Tab("🖼️ Image Generation"):
            img_prompt = gr.Textbox(label="Prompt")
            img_model = gr.Dropdown(choices=list(image_models.keys()), value="Stable Diffusion 1.5 (light)", label="Image Model")
            img_seed = gr.Slider(0, MAX_SEED, value=42, label="Seed")
            img_rand = gr.Checkbox(label="Randomize seed", value=True)
            img_btn = gr.Button("Generate Image")
            img_out = gr.Image()
            img_btn.click(fn=generate_image, inputs=[img_prompt, img_model, img_seed, img_rand], outputs=[img_out, img_seed])

        # Text Generation
        with gr.Tab("📝 Text Generation"):
            txt_prompt = gr.Textbox(label="Prompt")
            txt_model = gr.Dropdown(choices=list(text_models.keys()), value="GPT-2 (light)", label="Text Model")
            txt_btn = gr.Button("Generate Text")
            txt_out = gr.Textbox(label="Output Text")
            txt_btn.click(fn=generate_text, inputs=[txt_prompt, txt_model], outputs=[txt_out])

        # Video Generation
        with gr.Tab("🎥 Video Generation"):
            vid_prompt = gr.Textbox(label="Prompt")
            vid_model = gr.Dropdown(choices=list(video_models.keys()), value="CogVideoX-2B", label="Video Model")
            vid_seed = gr.Slider(0, MAX_SEED, value=42, label="Seed")
            vid_rand = gr.Checkbox(label="Randomize seed", value=True)
            vid_btn = gr.Button("Generate Video")
            vid_out = gr.Video()
            vid_btn.click(fn=generate_video, inputs=[vid_prompt, vid_model, vid_seed, vid_rand], outputs=[vid_out, vid_seed])

demo.launch(show_error=True)