models.py

import torch
from diffusers import StableDiffusionPipeline
from transformers import pipeline
import spaces
from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig

# Initialize models
sd_pipe = None
audio_pipe = None

@spaces.GPU(duration=1500)
def load_models():
    global sd_pipe, audio_pipe
    
    # Load Stable Diffusion for image generation
    sd_pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-1",
        torch_dtype=torch.float16
    ).to("cuda")
    
    # Apply AoT compilation to the transformer
    @spaces.GPU(duration=1500)
    def compile_transformer():
        with spaces.aoti_capture(sd_pipe.unet) as call:
            sd_pipe("AOT compilation", num_inference_steps=1)
        exported = torch.export.export(
            sd_pipe.unet,
            args=call.args,
            kwargs=call.kwargs
        )
        return spaces.aoti_compile(exported)
    
    compiled_transformer = compile_transformer()
    spaces.aoti_apply(compiled_transformer, sd_pipe.unet)
    
    # Apply FP8 quantization for additional speedup
    quantize_(sd_pipe.unet, Float8DynamicActivationFloat8WeightConfig())
    
    # Load MusicGen for audio generation
    audio_pipe = pipeline(
        "text-to-audio", 
        model="facebook/musicgen-small",
        device="cuda",
        torch_dtype=torch.float16
    )

@spaces.GPU(duration=120)
def generate_image(prompt: str) -> "PIL.Image":
    if sd_pipe is None:
        load_models()
    return sd_pipe(prompt, num_inference_steps=25).images[0]

@spaces.GPU(duration=60)
def generate_audio(prompt: str, duration: float) -> tuple:
    if audio_pipe is None:
        load_models()
    output = audio_pipe(
        prompt,
        forward_params={"max_new_tokens": int(duration * 50)},
    )
    return output["sampling_rate"], output["audio"][0]