xhs_factory/sd_service.py

"""
Stable Diffusion 服务模块
封装对 SD WebUI API 的调用，支持 txt2img 和 img2img
"""
import requests
import base64
import io
import logging
from PIL import Image

logger = logging.getLogger(__name__)

SD_TIMEOUT = 900  # 图片生成可能需要较长时间

# 默认反向提示词（针对 JuggernautXL / SDXL 优化）
DEFAULT_NEGATIVE = (
    "nsfw, nudity, lowres, bad anatomy, bad hands, text, error, missing fingers, "
    "extra digit, fewer digits, cropped, worst quality, low quality, normal quality, "
    "jpeg artifacts, signature, watermark, blurry, deformed, mutated, disfigured, "
    "ugly, duplicate, morbid, mutilated, poorly drawn face, poorly drawn hands, "
    "extra limbs, fused fingers, too many fingers, long neck, username, "
    "out of frame, distorted, oversaturated, underexposed, overexposed"
)


class SDService:
    """Stable Diffusion WebUI API 封装"""

    def __init__(self, sd_url: str = "http://127.0.0.1:7860"):
        self.sd_url = sd_url.rstrip("/")

    def check_connection(self) -> tuple[bool, str]:
        """检查 SD 服务是否可用"""
        try:
            resp = requests.get(f"{self.sd_url}/sdapi/v1/sd-models", timeout=5)
            if resp.status_code == 200:
                count = len(resp.json())
                return True, f"SD 已连接，{count} 个模型可用"
            return False, f"SD 返回异常状态: {resp.status_code}"
        except requests.exceptions.ConnectionError:
            return False, "SD WebUI 未启动或端口错误"
        except Exception as e:
            return False, f"SD 连接失败: {e}"

    def get_models(self) -> list[str]:
        """获取 SD 模型列表"""
        resp = requests.get(f"{self.sd_url}/sdapi/v1/sd-models", timeout=5)
        resp.raise_for_status()
        return [m["title"] for m in resp.json()]

    def switch_model(self, model_name: str):
        """切换 SD 模型"""
        try:
            requests.post(
                f"{self.sd_url}/sdapi/v1/options",
                json={"sd_model_checkpoint": model_name},
                timeout=60,
            )
        except Exception as e:
            logger.warning("模型切换失败: %s", e)

    def txt2img(
        self,
        prompt: str,
        negative_prompt: str = DEFAULT_NEGATIVE,
        model: str = None,
        steps: int = 30,
        cfg_scale: float = 5.0,
        width: int = 832,
        height: int = 1216,
        batch_size: int = 2,
        seed: int = -1,
        sampler_name: str = "DPM++ 2M",
        scheduler: str = "Karras",
    ) -> list[Image.Image]:
        """文生图（参数针对 JuggernautXL 优化）"""
        if model:
            self.switch_model(model)

        payload = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "steps": steps,
            "cfg_scale": cfg_scale,
            "width": width,
            "height": height,
            "batch_size": batch_size,
            "seed": seed,
            "sampler_name": sampler_name,
            "scheduler": scheduler,
        }

        resp = requests.post(
            f"{self.sd_url}/sdapi/v1/txt2img",
            json=payload,
            timeout=SD_TIMEOUT,
        )
        resp.raise_for_status()

        images = []
        for img_b64 in resp.json().get("images", []):
            img = Image.open(io.BytesIO(base64.b64decode(img_b64)))
            images.append(img)
        return images

    def img2img(
        self,
        init_image: Image.Image,
        prompt: str,
        negative_prompt: str = DEFAULT_NEGATIVE,
        denoising_strength: float = 0.5,
        steps: int = 30,
        cfg_scale: float = 5.0,
        sampler_name: str = "DPM++ 2M",
        scheduler: str = "Karras",
    ) -> list[Image.Image]:
        """图生图（参数针对 JuggernautXL 优化）"""
        # 将 PIL Image 转为 base64
        buf = io.BytesIO()
        init_image.save(buf, format="PNG")
        init_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        payload = {
            "init_images": [init_b64],
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "denoising_strength": denoising_strength,
            "steps": steps,
            "cfg_scale": cfg_scale,
            "width": init_image.width,
            "height": init_image.height,
            "sampler_name": sampler_name,
            "scheduler": scheduler,
        }

        resp = requests.post(
            f"{self.sd_url}/sdapi/v1/img2img",
            json=payload,
            timeout=SD_TIMEOUT,
        )
        resp.raise_for_status()

        images = []
        for img_b64 in resp.json().get("images", []):
            img = Image.open(io.BytesIO(base64.b64decode(img_b64)))
            images.append(img)
        return images

    def get_lora_models(self) -> list[str]:
        """获取可用的 LoRA 模型列表"""
        try:
            resp = requests.get(f"{self.sd_url}/sdapi/v1/loras", timeout=5)
            resp.raise_for_status()
            return [lora["name"] for lora in resp.json()]
        except Exception:
            return []