zhoujie
358b957f5d
- 新增模型降级机制，当主模型失败时自动尝试备选模型列表【FALLBACK_MODELS】 - 增强 `_chat` 方法，支持空返回检测、json_mode 回退和多重错误处理 - 重构 `_parse_json` 方法，实现五重容错解析策略以应对不同模型的输出格式 - 为 `generate_copy`、`generate_copy_with_reference` 和 `analyze_hotspots` 方法添加重试逻辑，在 JSON 解析失败时自动关闭 json_mode 重试 🔧 chore(config): 更新默认模型配置与安全令牌 - 将默认 LLM 模型从 `gemini-3-flash-preview` 更改为 `deepseek-v3` - 更新 `xsec_token` 安全令牌 ✨ feat(sd): 集成 ReActor 换脸功能并扩展人设主题池 - 在 `SDService` 中新增头像管理静态方法 (`load_face_image`, `save_face_image`) 和 ReActor 参数构建方法 - 为 `txt2img` 方法添加 `face_image` 参数，支持在生成图片时自动换脸 - 在 `main.py` 的 Web UI 中新增头像上传、预览与管理界面 - 扩展 `generate_images` 函数，支持根据复选框状态启用换脸功能 - 重构人设系统，为 24 种预设人设分别定义专属的【主题池】和【评论关键词池】，并实现人设切换时的自动联动更新 - 在自动化发布 (`auto_publish_once`) 和定时调度 (`_scheduler_loop`) 中集成换脸选项 📝 docs(main): 添加新图片资源 - 新增图片资源文件：`beauty.png`, `my_face.png`, `myself.jpg`, `zjz.png`
267 lines
10 KiB
Python
267 lines
10 KiB
Python
"""
|
||
Stable Diffusion 服务模块
|
||
封装对 SD WebUI API 的调用，支持 txt2img 和 img2img，支持 ReActor 换脸
|
||
"""
|
||
import requests
|
||
import base64
|
||
import io
|
||
import logging
|
||
import os
|
||
from PIL import Image
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
SD_TIMEOUT = 1800 # 图片生成可能需要较长时间
|
||
|
||
# 头像文件默认保存路径
|
||
FACE_IMAGE_PATH = os.path.join(os.path.dirname(__file__), "my_face.png")
|
||
|
||
# 默认反向提示词（针对 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, "
|
||
"western face, european face, caucasian, deep-set eyes, high nose bridge, "
|
||
"blonde hair, red hair, blue eyes, green eyes, freckles, thick body hair"
|
||
)
|
||
|
||
|
||
class SDService:
|
||
"""Stable Diffusion WebUI API 封装"""
|
||
|
||
def __init__(self, sd_url: str = "http://127.0.0.1:7860"):
|
||
self.sd_url = sd_url.rstrip("/")
|
||
|
||
# ---------- 工具方法 ----------
|
||
|
||
@staticmethod
|
||
def _image_to_base64(img: Image.Image) -> str:
|
||
"""PIL Image → base64 字符串"""
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||
|
||
@staticmethod
|
||
def load_face_image(path: str = None) -> Image.Image | None:
|
||
"""加载头像图片，不存在则返回 None"""
|
||
path = path or FACE_IMAGE_PATH
|
||
if path and os.path.isfile(path):
|
||
try:
|
||
return Image.open(path).convert("RGB")
|
||
except Exception as e:
|
||
logger.warning("头像加载失败: %s", e)
|
||
return None
|
||
|
||
@staticmethod
|
||
def save_face_image(img: Image.Image, path: str = None) -> str:
|
||
"""保存头像图片，返回保存路径"""
|
||
path = path or FACE_IMAGE_PATH
|
||
img = img.convert("RGB")
|
||
img.save(path, format="PNG")
|
||
logger.info("头像已保存: %s", path)
|
||
return path
|
||
|
||
def _build_reactor_args(self, face_image: Image.Image) -> dict:
|
||
"""构建 ReActor 换脸参数（alwayson_scripts 格式）
|
||
|
||
参数索引对照 (reactor script-info):
|
||
0: source_image (base64) 1: enable 2: source_faces
|
||
3: target_faces 4: model 5: restore_face
|
||
6: restore_visibility 7: restore_first 8: upscaler
|
||
9: scale 10: upscaler_vis 11: swap_in_source
|
||
12: swap_in_generated 13: log_level 14: gender_source
|
||
15: gender_target 16: save_original 17: codeformer_weight
|
||
18: source_hash_check 19: target_hash_check 20: exec_provider
|
||
21: face_mask_correction 22: select_source 23: face_model
|
||
24: source_folder 25: multiple_sources 26: random_image
|
||
27: force_upscale 28: threshold 29: max_faces
|
||
30: tab_single
|
||
"""
|
||
face_b64 = self._image_to_base64(face_image)
|
||
return {
|
||
"reactor": {
|
||
"args": [
|
||
face_b64, # 0: source image (base64)
|
||
True, # 1: enable ReActor
|
||
"0", # 2: source face index
|
||
"0", # 3: target face index
|
||
"inswapper_128.onnx", # 4: swap model
|
||
"CodeFormer", # 5: restore face method
|
||
1, # 6: restore face visibility
|
||
True, # 7: restore face first, then upscale
|
||
"None", # 8: upscaler
|
||
1, # 9: scale
|
||
1, # 10: upscaler visibility
|
||
False, # 11: swap in source
|
||
True, # 12: swap in generated
|
||
"Minimum", # 13: log level
|
||
"No", # 14: gender detection (source)
|
||
"No", # 15: gender detection (target)
|
||
False, # 16: save original
|
||
0.6, # 17: CodeFormer weight (fidelity)
|
||
True, # 18: source hash check
|
||
False, # 19: target hash check
|
||
"CUDA", # 20: execution provider
|
||
True, # 21: face mask correction
|
||
"Image(s)", # 22: select source type
|
||
"None", # 23: face model name
|
||
"", # 24: source folder
|
||
None, # 25: multiple source images
|
||
False, # 26: random image
|
||
False, # 27: force upscale
|
||
0.5, # 28: detection threshold
|
||
0, # 29: max faces (0 = no limit)
|
||
"tab_single", # 30: tab
|
||
],
|
||
}
|
||
}
|
||
|
||
def has_reactor(self) -> bool:
|
||
"""检查 SD WebUI 是否安装了 ReActor 扩展"""
|
||
try:
|
||
resp = requests.get(f"{self.sd_url}/sdapi/v1/scripts", timeout=5)
|
||
scripts = resp.json()
|
||
all_scripts = scripts.get("txt2img", []) + scripts.get("img2img", [])
|
||
return any("reactor" in s.lower() for s in all_scripts)
|
||
except Exception:
|
||
return False
|
||
|
||
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",
|
||
face_image: Image.Image = None,
|
||
) -> list[Image.Image]:
|
||
"""文生图（参数针对 JuggernautXL 优化）
|
||
|
||
Args:
|
||
face_image: 头像 PIL Image，传入后自动启用 ReActor 换脸
|
||
"""
|
||
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,
|
||
}
|
||
|
||
# 如果提供了头像，通过 ReActor 换脸
|
||
if face_image is not None:
|
||
payload["alwayson_scripts"] = self._build_reactor_args(face_image)
|
||
logger.info("🎭 ReActor 换脸已启用")
|
||
|
||
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 []
|