File size: 1,508 Bytes
8ab53f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """
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"} |