Mortdecai
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
"""SDXL Turbo wrapper for horror image generation."""
|
|
|
|
import io
|
|
|
|
import torch
|
|
|
|
try:
|
|
from diffusers import AutoPipelineForText2Image
|
|
except ImportError: # pragma: no cover - exercised only in test environments
|
|
AutoPipelineForText2Image = None # Tests patch this attribute directly.
|
|
|
|
from server.config import config
|
|
from server.prompts import NEGATIVE_PROMPT
|
|
|
|
|
|
class AssetGenerator:
|
|
"""Generates horror images via SDXL Turbo."""
|
|
|
|
def __init__(self, device: str | None = None, model_id: str | None = None):
|
|
self.device = device or config.device
|
|
self.model_id = model_id or config.models.sdxl_model_id
|
|
|
|
if AutoPipelineForText2Image is None:
|
|
raise RuntimeError(
|
|
"diffusers is not installed; install diffusers to use AssetGenerator"
|
|
)
|
|
|
|
use_fp16 = self.device == "cuda"
|
|
self._pipe = AutoPipelineForText2Image.from_pretrained(
|
|
self.model_id,
|
|
torch_dtype=torch.float16 if use_fp16 else torch.float32,
|
|
variant="fp16" if use_fp16 else None,
|
|
)
|
|
if self.device == "cuda":
|
|
self._pipe = self._pipe.to("cuda")
|
|
|
|
def generate(self, prompt: str, seed: int | None = None) -> bytes:
|
|
"""Generate a 512x512 PNG image from a horror prompt. Returns PNG bytes."""
|
|
if seed is None:
|
|
seed = torch.randint(0, 2**32, (1,)).item()
|
|
|
|
generator = torch.Generator(device=self.device).manual_seed(seed)
|
|
|
|
result = self._pipe(
|
|
prompt=prompt,
|
|
negative_prompt=NEGATIVE_PROMPT,
|
|
num_inference_steps=config.models.sdxl_steps,
|
|
guidance_scale=config.models.sdxl_guidance_scale,
|
|
width=config.models.sdxl_width,
|
|
height=config.models.sdxl_height,
|
|
generator=generator,
|
|
)
|
|
|
|
image = result.images[0]
|
|
buf = io.BytesIO()
|
|
image.save(buf, format="PNG")
|
|
return buf.getvalue()
|