Stable Diffusion WebUI integration. File CORP proxy.
This commit is contained in:
@@ -0,0 +1,167 @@
|
||||
import requests
|
||||
import tomllib
|
||||
import os
|
||||
import base64
|
||||
import config
|
||||
from typing import Any, Dict, List
|
||||
from tools.utils import ToolError
|
||||
|
||||
# Paths to config files
|
||||
MODEL_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "model_presets.toml")
|
||||
RES_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resolution_presets.toml")
|
||||
|
||||
SD_URL = "http://127.0.0.1:7860"
|
||||
|
||||
def load_toml(path: str) -> Dict[str, Any]:
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
return tomllib.load(f)
|
||||
except Exception as e:
|
||||
raise ToolError(f"Failed to load config file {path}: {str(e)}")
|
||||
|
||||
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generates an image using the specified model and parameters.
|
||||
"""
|
||||
# 1. REQUIRED ARGUMENTS
|
||||
prompt = args.get("prompt")
|
||||
if not prompt:
|
||||
raise ToolError("The 'prompt' argument is required.")
|
||||
|
||||
model_name = args.get("model_name")
|
||||
if not model_name:
|
||||
raise ToolError("The 'model_name' argument is required. Use get_model_info to see available models.")
|
||||
|
||||
# 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
|
||||
try:
|
||||
info_resp = requests.get(f"{SD_URL}/info", timeout=10)
|
||||
info_resp.raise_for_status()
|
||||
info_data = info_resp.json()
|
||||
params_info = info_data["named_endpoints"]["/txt2img"]["parameters"]
|
||||
except Exception as e:
|
||||
raise ToolError(f"Failed to connect to Stable Diffusion server: {str(e)}")
|
||||
|
||||
# Build the label-to-index map
|
||||
label_map = {}
|
||||
label_counts = {}
|
||||
|
||||
payload = [p["parameter_default"] for p in params_info]
|
||||
|
||||
for idx, p in enumerate(params_info):
|
||||
label = p["label"]
|
||||
if not label or label.startswith("parameter_"):
|
||||
label = p["parameter_name"]
|
||||
|
||||
if label in label_counts:
|
||||
label_counts[label] += 1
|
||||
mapped_label = f"{label} [{label_counts[label]}]"
|
||||
else:
|
||||
label_counts[label] = 0
|
||||
mapped_label = label
|
||||
|
||||
label_map[mapped_label] = idx
|
||||
|
||||
# 3. LOAD CONFIGS
|
||||
models_cfg = load_toml(MODEL_PRESETS_PATH)
|
||||
res_cfg = load_toml(RES_PRESETS_PATH)
|
||||
|
||||
if model_name not in models_cfg:
|
||||
raise ToolError(f"Model '{model_name}' not found in presets. Available: {', '.join(models_cfg.keys())}")
|
||||
|
||||
model_preset = models_cfg[model_name]
|
||||
|
||||
# 4. MERGE PIPELINE
|
||||
for key, value in model_preset.items():
|
||||
if key in ["description", "guide", "Resolution Set"]:
|
||||
continue
|
||||
if key in label_map:
|
||||
payload[label_map[key]] = value
|
||||
elif key.startswith("param_"):
|
||||
try:
|
||||
idx = int(key.replace("param_", ""))
|
||||
if 0 <= idx < len(payload):
|
||||
payload[idx] = value
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
res_preset_name = args.get("resolution_preset")
|
||||
if res_preset_name:
|
||||
res_set_name = model_preset.get("Resolution Set")
|
||||
if res_set_name and res_set_name in res_cfg:
|
||||
res_set = res_cfg[res_set_name]
|
||||
if res_preset_name in res_set:
|
||||
res_vals = res_set[res_preset_name]
|
||||
if "Width" in label_map:
|
||||
payload[label_map["Width"]] = res_vals["width"]
|
||||
if "Height" in label_map:
|
||||
payload[label_map["Height"]] = res_vals["height"]
|
||||
else:
|
||||
raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. Available: {', '.join(res_set.keys())}")
|
||||
else:
|
||||
raise ToolError(f"No resolution set configured for model '{model_name}'.")
|
||||
|
||||
overrides = {
|
||||
"prompt": "Prompt",
|
||||
"negative_prompt": "Negative Prompt",
|
||||
"steps": "Sampling Steps",
|
||||
"cfg_scale": "CFG Scale",
|
||||
"seed": "Seed",
|
||||
"sampler": "Sampling Method"
|
||||
}
|
||||
|
||||
for arg_key, label in overrides.items():
|
||||
if arg_key in args and label in label_map:
|
||||
payload[label_map[label]] = args[arg_key]
|
||||
|
||||
if "Prompt" in label_map:
|
||||
payload[label_map["Prompt"]] = prompt
|
||||
|
||||
for _ in range(7):
|
||||
payload.insert(39, None)
|
||||
|
||||
# 6. EXECUTE GENERATION
|
||||
try:
|
||||
gen_payload = {"data": payload}
|
||||
gen_resp = requests.post(f"{SD_URL}/api/txt2img", json=gen_payload, timeout=300)
|
||||
gen_resp.raise_for_status()
|
||||
res_data = gen_resp.json()
|
||||
|
||||
gallery_data = res_data["data"][0]["value"]
|
||||
if not gallery_data:
|
||||
raise ToolError("Server returned successfully but no image was generated.")
|
||||
|
||||
sd_img_url = gallery_data[0]["image"]["url"]
|
||||
|
||||
# Extract raw path from SD URL (e.g. http://.../file=/tmp/gradio/abc.png -> /tmp/gradio/abc.png)
|
||||
if "/file=" in sd_img_url:
|
||||
raw_path = sd_img_url.split("/file=")[1]
|
||||
else:
|
||||
# Fallback if the URL format changes
|
||||
raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")
|
||||
|
||||
# Construct proxy URL through our MCP server
|
||||
# We strip leading slash from raw_path to avoid double slashes in the proxy URL if we want,
|
||||
# but the current endpoint handles it.
|
||||
proxy_url = f"http://localhost:{config.PORT}/file{raw_path}"
|
||||
|
||||
# Download the image and convert to base64 for the AI's vision
|
||||
img_resp = requests.get(sd_img_url, timeout=30)
|
||||
img_resp.raise_for_status()
|
||||
b64_data = base64.b64encode(img_resp.content).decode('utf-8')
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"Image successfully generated using model '{model_name}'.\n\nTo show the image to the user, you can include this markdown link in your response: "
|
||||
},
|
||||
{
|
||||
"type": "image",
|
||||
"data": b64_data,
|
||||
"mimeType": "image/png"
|
||||
}
|
||||
]
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
raise ToolError(f"Server error during generation: {str(e)}")
|
||||
except Exception as e:
|
||||
raise ToolError(f"Unexpected error during generation: {str(e)}")
|
||||
Reference in New Issue
Block a user