| """ |
| 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"} |
