zhoujie
88faca150d
- 新增项目配置文件(.gitignore, config.json)和核心文档(Todo.md, mcp.md) - 实现配置管理模块(config_manager.py)，支持单例模式和自动保存 - 实现LLM服务模块(llm_service.py)，包含文案生成、热点分析、评论回复等Prompt模板 - 实现SD服务模块(sd_service.py)，封装Stable Diffusion WebUI API调用 - 实现MCP客户端模块(mcp_client.py)，封装小红书MCP服务HTTP调用 - 实现主程序(main.py)，构建Gradio界面，包含内容创作、热点探测、评论管家、账号登录、数据看板五大功能模块 - 保留V1版本备份(main_v1_backup.py)供参考 - 添加项目依赖文件(requirements.txt)
146 lines
4.4 KiB
Python
146 lines
4.4 KiB
Python
"""
|
|
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 = 180 # 图片生成可能需要较长时间
|
|
|
|
# 默认反向提示词
|
|
DEFAULT_NEGATIVE = (
|
|
"nsfw, 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"
|
|
)
|
|
|
|
|
|
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 = 25,
|
|
cfg_scale: float = 7.0,
|
|
width: int = 768,
|
|
height: int = 1024,
|
|
batch_size: int = 2,
|
|
seed: int = -1,
|
|
) -> list[Image.Image]:
|
|
"""文生图"""
|
|
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,
|
|
}
|
|
|
|
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.6,
|
|
steps: int = 25,
|
|
cfg_scale: float = 7.0,
|
|
) -> list[Image.Image]:
|
|
"""图生图（参考图修改）"""
|
|
# 将 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,
|
|
}
|
|
|
|
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 []
|