"""
Unified Model Adapter & Native Diffusers Driver
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""

from abc import ABC, abstractmethod
import torch
from typing import Dict, Any, Optional

class AbstractModelAdapter(ABC):
    @abstractmethod
    async def load_model(self, model_path: str, device: str, quantization: Optional[str] = None):
        pass

    @abstractmethod
    async def generate(self, prompt: str, negative_prompt: str, **kwargs) -> Dict[str, Any]:
        pass

class DiffusersVideoAdapter(AbstractModelAdapter):
    """Native Diffusers 视频模型适配器"""
    def __init__(self):
        self.pipeline = None
        self.device = "cuda"

    async def load_model(self, model_path: str, device: str = "cuda", quantization: Optional[str] = "fp8"):
        self.device = device
        dtype = torch.float8_e4m3fn if quantization == "fp8" else torch.float16
        from diffusers import DiffusionPipeline
        self.pipeline = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True)
        if hasattr(self.pipeline, "enable_model_cpu_offload"):
            self.pipeline.enable_model_cpu_offload()

    async def generate(self, prompt: str, negative_prompt: str = "", **kwargs) -> Dict[str, Any]:
        output = self.pipeline(prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator(device=self.device).manual_seed(42))
        return {"video_frames": output.frames[0], "status": "SUCCESS"}