""" 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") # ==================== 多模型配置系统 ==================== # 每个模型的最优参数、prompt 增强词、负面提示词、三档预设 SD_MODEL_PROFILES = { # ---- majicmixRealistic: 东亚网红感，朋友圈自拍/美妆/穿搭 (SD 1.5) ---- "majicmixRealistic": { "display_name": "majicmixRealistic ⭐⭐⭐⭐⭐", "description": "东亚网红感 | 朋友圈自拍、美妆、穿搭", "arch": "sd15", # SD 1.5 架构 # 自动追加到 prompt 前面的增强词 "prompt_prefix": ( "(best quality:1.4), (masterpiece:1.4), (ultra detailed:1.3), " "(photorealistic:1.4), (realistic:1.3), raw photo, " "(asian girl:1.3), (chinese:1.2), (east asian features:1.2), " "(delicate facial features:1.2), (fair skin:1.1), (natural skin texture:1.2), " "(soft lighting:1.1), (natural makeup:1.1), " ), # 自动追加到 prompt 后面的补充词 "prompt_suffix": ( ", film grain, shallow depth of field, " "instagram aesthetic, xiaohongshu style, phone camera feel" ), "negative_prompt": ( "(nsfw:1.5), (nudity:1.5), (worst quality:2), (low quality:2), (normal quality:2), " "lowres, bad anatomy, bad hands, text, error, missing fingers, " "extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, " "blurry, deformed, mutated, disfigured, ugly, duplicate, " "poorly drawn face, poorly drawn hands, extra limbs, fused fingers, " "too many fingers, long neck, out of frame, " "western face, european face, caucasian, deep-set eyes, high nose bridge, " "blonde hair, red hair, blue eyes, green eyes, freckles, thick body hair, " "painting, cartoon, anime, sketch, illustration, 3d render" ), "presets": { "快速 (约30秒)": { "steps": 20, "cfg_scale": 7.0, "width": 512, "height": 768, "sampler_name": "Euler a", "scheduler": "Normal", "batch_size": 2, }, "标准 (约1分钟)": { "steps": 30, "cfg_scale": 7.0, "width": 512, "height": 768, "sampler_name": "DPM++ 2M", "scheduler": "Karras", "batch_size": 2, }, "精细 (约2-3分钟)": { "steps": 40, "cfg_scale": 7.5, "width": 576, "height": 864, "sampler_name": "DPM++ SDE", "scheduler": "Karras", "batch_size": 2, }, }, }, # ---- Realistic Vision: 写实摄影感，纪实摄影/街拍/真实质感 (SD 1.5) ---- "realisticVision": { "display_name": "Realistic Vision ⭐⭐⭐⭐", "description": "写实摄影感 | 纪实摄影、街拍、真实质感", "arch": "sd15", "prompt_prefix": ( "RAW photo, (best quality:1.4), (masterpiece:1.3), (realistic:1.4), " "(photorealistic:1.4), 8k uhd, DSLR, high quality, " "(asian:1.2), (chinese girl:1.2), (east asian features:1.1), " "(natural skin:1.2), (skin pores:1.1), (detailed skin texture:1.2), " ), "prompt_suffix": ( ", shot on Canon EOS R5, 85mm lens, f/1.8, " "natural lighting, documentary style, street photography, " "film color grading, depth of field" ), "negative_prompt": ( "(nsfw:1.5), (nudity:1.5), (worst quality:2), (low quality:2), (normal quality:2), " "lowres, bad anatomy, bad hands, text, error, missing fingers, " "extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, " "blurry, deformed, mutated, disfigured, ugly, duplicate, " "poorly drawn face, extra limbs, fused fingers, long neck, " "western face, european face, caucasian, deep-set eyes, " "blonde hair, blue eyes, green eyes, freckles, " "painting, cartoon, anime, sketch, illustration, 3d render, " "over-sharpened, over-saturated, plastic skin, airbrushed, " "smooth skin, doll-like, HDR, overprocessed" ), "presets": { "快速 (约30秒)": { "steps": 20, "cfg_scale": 7.0, "width": 512, "height": 768, "sampler_name": "Euler a", "scheduler": "Normal", "batch_size": 2, }, "标准 (约1分钟)": { "steps": 28, "cfg_scale": 7.0, "width": 512, "height": 768, "sampler_name": "DPM++ 2M", "scheduler": "Karras", "batch_size": 2, }, "精细 (约2-3分钟)": { "steps": 40, "cfg_scale": 7.5, "width": 576, "height": 864, "sampler_name": "DPM++ SDE", "scheduler": "Karras", "batch_size": 2, }, }, }, # ---- Juggernaut XL: 电影大片感，高画质/商业摄影/复杂背景 (SDXL) ---- "juggernautXL": { "display_name": "Juggernaut XL ⭐⭐⭐⭐", "description": "电影大片感 | 高画质、商业摄影、复杂背景", "arch": "sdxl", # SDXL 架构 "prompt_prefix": ( "masterpiece, best quality, ultra detailed, 8k uhd, high resolution, " "photorealistic, cinematic lighting, cinematic composition, " "asian girl, chinese, east asian features, black hair, dark brown eyes, " "delicate facial features, fair skin, slim figure, " ), "prompt_suffix": ( ", cinematic color grading, anamorphic lens, bokeh, " "volumetric lighting, ray tracing, global illumination, " "commercial photography, editorial style, vogue aesthetic" ), "negative_prompt": ( "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" ), "presets": { "快速 (约30秒)": { "steps": 12, "cfg_scale": 5.0, "width": 768, "height": 1024, "sampler_name": "Euler a", "scheduler": "Normal", "batch_size": 2, }, "标准 (约1分钟)": { "steps": 20, "cfg_scale": 5.5, "width": 832, "height": 1216, "sampler_name": "DPM++ 2M", "scheduler": "Karras", "batch_size": 2, }, "精细 (约2-3分钟)": { "steps": 35, "cfg_scale": 6.0, "width": 832, "height": 1216, "sampler_name": "DPM++ 2M SDE", "scheduler": "Karras", "batch_size": 2, }, }, }, } # 默认配置 profile key DEFAULT_MODEL_PROFILE = "juggernautXL" def detect_model_profile(model_name: str) -> str: """根据 SD 模型名称自动识别对应的 profile key""" name_lower = model_name.lower() if model_name else "" if "majicmix" in name_lower or "majic" in name_lower: return "majicmixRealistic" elif "realistic" in name_lower and "vision" in name_lower: return "realisticVision" elif "rv" in name_lower and ("v5" in name_lower or "v6" in name_lower or "v4" in name_lower): return "realisticVision" # RV v5.1 等简写 elif "juggernaut" in name_lower or "jugger" in name_lower: return "juggernautXL" # 根据架构猜测 elif "xl" in name_lower or "sdxl" in name_lower: return "juggernautXL" # SDXL 架构默认用 Juggernaut 参数 else: return DEFAULT_MODEL_PROFILE # 无法识别时默认 def get_model_profile(model_name: str = None) -> dict: """获取模型配置 profile""" key = detect_model_profile(model_name) if model_name else DEFAULT_MODEL_PROFILE return SD_MODEL_PROFILES.get(key, SD_MODEL_PROFILES[DEFAULT_MODEL_PROFILE]) def get_model_profile_info(model_name: str = None) -> str: """获取当前模型的显示信息 (Markdown 格式)""" profile = get_model_profile(model_name) key = detect_model_profile(model_name) if model_name else DEFAULT_MODEL_PROFILE is_default = key == DEFAULT_MODEL_PROFILE and model_name and detect_model_profile(model_name) == DEFAULT_MODEL_PROFILE # 如果检测结果是默认回退的, 说明是未知模型 actual_key = detect_model_profile(model_name) if model_name else None presets = profile["presets"] first_preset = list(presets.values())[0] res = f"{first_preset.get('width', '?')}×{first_preset.get('height', '?')}" lines = [ f"**🎨 {profile['display_name']}** | `{profile['arch'].upper()}` | {res}", f"> {profile['description']}", ] if model_name and not any(k in (model_name or "").lower() for k in ["majicmix", "realistic", "juggernaut"]): lines.append(f"> ⚠️ 未识别的模型，使用默认档案 ({profile['display_name']})") return "\n".join(lines) # ==================== 兼容旧接口 ==================== # 默认预设和反向提示词 (使用 Juggernaut XL 作为默认) SD_PRESETS = SD_MODEL_PROFILES[DEFAULT_MODEL_PROFILE]["presets"] SD_PRESET_NAMES = list(SD_PRESETS.keys()) def get_sd_preset(name: str, model_name: str = None) -> dict: """获取生成预设参数，自动适配模型""" profile = get_model_profile(model_name) presets = profile.get("presets", SD_PRESETS) return presets.get(name, presets.get("标准 (约1分钟)", list(presets.values())[0])) # 默认反向提示词（Juggernaut XL） DEFAULT_NEGATIVE = SD_MODEL_PROFILES[DEFAULT_MODEL_PROFILE]["negative_prompt"] 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 1, # 13: console log level (0=min, 1=med, 2=max) 0, # 14: gender detection source (0=No) 0, # 15: gender detection target (0=No) False, # 16: save original 0.8, # 17: CodeFormer weight (0=max effect, 1=min) False, # 18: source hash check False, # 19: target hash check "CUDA", # 20: execution provider True, # 21: face mask correction 0, # 22: select source (0=Image, 1=FaceModel, 2=Folder) "", # 23: face model filename (when #22=1) "", # 24: source folder path (when #22=2) None, # 25: skip for API False, # 26: random image False, # 27: force upscale 0.6, # 28: face detection threshold 2, # 29: max faces to detect (0=unlimited) ], } } 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 = None, model: str = None, steps: int = None, cfg_scale: float = None, width: int = None, height: int = None, batch_size: int = None, seed: int = -1, sampler_name: str = None, scheduler: str = None, face_image: Image.Image = None, quality_mode: str = None, ) -> list[Image.Image]: """文生图（自动适配当前 SD 模型的最优参数） Args: model: SD 模型名，自动识别并应用对应配置 face_image: 头像 PIL Image，传入后自动启用 ReActor 换脸 quality_mode: 预设模式名 """ if model: self.switch_model(model) # 自动识别模型配置 profile = get_model_profile(model) profile_key = detect_model_profile(model) logger.info("🎯 SD 模型识别: %s → %s (%s)", model or "默认", profile_key, profile["description"]) # 加载模型专属预设参数 preset = get_sd_preset(quality_mode, model) if quality_mode else get_sd_preset("标准 (约1分钟)", model) # 自动增强 prompt: 前缀 + 原始 prompt + 后缀 enhanced_prompt = profile.get("prompt_prefix", "") + prompt + profile.get("prompt_suffix", "") # 使用模型专属反向提示词 final_negative = negative_prompt if negative_prompt is not None else profile.get("negative_prompt", DEFAULT_NEGATIVE) payload = { "prompt": enhanced_prompt, "negative_prompt": final_negative, "steps": steps if steps is not None else preset["steps"], "cfg_scale": cfg_scale if cfg_scale is not None else preset["cfg_scale"], "width": width if width is not None else preset["width"], "height": height if height is not None else preset["height"], "batch_size": batch_size if batch_size is not None else preset["batch_size"], "seed": seed, "sampler_name": sampler_name if sampler_name is not None else preset["sampler_name"], "scheduler": scheduler if scheduler is not None else preset["scheduler"], } logger.info("SD 生成参数 [%s]: steps=%s, cfg=%.1f, %dx%d, sampler=%s", profile_key, payload['steps'], payload['cfg_scale'], payload['width'], payload['height'], payload['sampler_name']) # 如果提供了头像，通过 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 = None, denoising_strength: float = 0.5, steps: int = 30, cfg_scale: float = None, sampler_name: str = None, scheduler: str = None, model: str = None, ) -> list[Image.Image]: """图生图（自动适配模型参数）""" profile = get_model_profile(model) preset = get_sd_preset("标准 (约1分钟)", model) # 将 PIL Image 转为 base64 buf = io.BytesIO() init_image.save(buf, format="PNG") init_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") enhanced_prompt = profile.get("prompt_prefix", "") + prompt + profile.get("prompt_suffix", "") final_negative = negative_prompt if negative_prompt is not None else profile.get("negative_prompt", DEFAULT_NEGATIVE) payload = { "init_images": [init_b64], "prompt": enhanced_prompt, "negative_prompt": final_negative, "denoising_strength": denoising_strength, "steps": steps, "cfg_scale": cfg_scale if cfg_scale is not None else preset["cfg_scale"], "width": init_image.width, "height": init_image.height, "sampler_name": sampler_name if sampler_name is not None else preset["sampler_name"], "scheduler": scheduler if scheduler is not None else preset["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 []