Fixed RPG

This commit is contained in:
2026-06-01 07:44:38 +03:00
parent 600ad78f05
commit d4cd8f02f4
30 changed files with 1516 additions and 816 deletions
+56 -22
View File
@@ -16,13 +16,26 @@ SD_STEPS = int(os.getenv("SD_STEPS", "28"))
SD_CFG = float(os.getenv("SD_CFG", "7"))
SD_SAMPLER = os.getenv("SD_SAMPLER", "euler")
SD_SCHEDULER = os.getenv("SD_SCHEDULER", "normal")
SD_CHECKPOINT = os.getenv("SD_CHECKPOINT", "NetaYumev35_pretrained_all_in_one.safetensors")
SD_CHECKPOINT = os.getenv("SD_CHECKPOINT", "")
SD_DEFAULT_NEGATIVE = os.getenv(
"SD_DEFAULT_NEGATIVE",
"low quality, worst quality, blurry, bad anatomy, watermark, text",
)
# Anima split-model settings
SD_UNET = os.getenv("SD_UNET", "anima-preview3-base.safetensors")
SD_CLIP = os.getenv("SD_CLIP", "qwen_3_06b_base.safetensors")
SD_VAE = os.getenv("SD_VAE", "qwen_image_vae.safetensors")
IMAGES_DIR = Path(os.getenv("IMAGES_DIR", "static/images"))
ANIMA_CHECKPOINTS = {"anima-preview3-base.safetensors"}
PONY_CHECKPOINTS = {"ponyDiffusionV6XL_v6StartWithThisOne.safetensors"}
def _use_anima() -> bool:
return bool(SD_UNET) and not SD_CHECKPOINT
def split_prompt_and_negative(full_prompt: str) -> tuple[str, str]:
if "\n\nNegative prompt:" in full_prompt:
@@ -32,26 +45,44 @@ def split_prompt_and_negative(full_prompt: str) -> tuple[str, str]:
def _build_workflow(positive: str, negative: str) -> dict:
"""Minimal KSampler workflow for ComfyUI API."""
seed = int(uuid.uuid4().int % 2**32)
if _use_anima():
return {
"44": {"class_type": "UNETLoader", "inputs": {"unet_name": SD_UNET, "weight_dtype": "default"}},
"45": {"class_type": "CLIPLoader", "inputs": {"clip_name": SD_CLIP, "type": "stable_diffusion", "device": "default"}},
"15": {"class_type": "VAELoader", "inputs": {"vae_name": SD_VAE}},
"28": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
"11": {"class_type": "CLIPTextEncode", "inputs": {"text": positive, "clip": ["45", 0]}},
"12": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["45", 0]}},
"19": {
"class_type": "KSampler",
"inputs": {
"model": ["44", 0], "positive": ["11", 0], "negative": ["12", 0],
"latent_image": ["28", 0], "seed": seed,
"steps": SD_STEPS, "cfg": SD_CFG,
"sampler_name": os.getenv("SD_SAMPLER", "er_sde"),
"scheduler": os.getenv("SD_SCHEDULER", "simple"),
"denoise": 1.0,
},
},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["19", 0], "vae": ["15", 0]}},
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "chatbot", "images": ["8", 0]}},
}
# Standard checkpoint workflow (Pony / SDXL)
return {
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": SD_CHECKPOINT}},
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 832, "height": 1216, "batch_size": 1}},
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": positive, "clip": ["4", 1]}},
"7": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["4", 1]}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["10", 0], "vae": ["4", 2]}},
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "chatbot", "images": ["8", 0]}},
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": SD_CHECKPOINT}},
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 832, "height": 1216, "batch_size": 1}},
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": positive, "clip": ["4", 1]}},
"7": {"class_type": "CLIPTextEncode", "inputs": {"text": negative, "clip": ["4", 1]}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["10", 0], "vae": ["4", 2]}},
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "chatbot", "images": ["8", 0]}},
"10": {
"class_type": "KSampler",
"inputs": {
"model": ["4", 0],
"positive": ["6", 0],
"negative": ["7", 0],
"latent_image": ["5", 0],
"seed": int(uuid.uuid4().int % 2**32),
"steps": SD_STEPS,
"cfg": SD_CFG,
"sampler_name": SD_SAMPLER,
"scheduler": SD_SCHEDULER,
"model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0],
"latent_image": ["5", 0], "seed": seed,
"steps": SD_STEPS, "cfg": SD_CFG,
"sampler_name": SD_SAMPLER, "scheduler": SD_SCHEDULER,
"denoise": 1.0,
},
},
@@ -74,7 +105,6 @@ async def txt2img(prompt: str, negative_prompt: str | None = None) -> tuple[byte
logger.info("ComfyUI request → %s prompt: %.120s", SD_BASE_URL, prompt)
async with httpx.AsyncClient(timeout=300) as client:
# queue the prompt
resp = await client.post(
f"{SD_BASE_URL}/prompt",
json={"prompt": workflow, "client_id": client_id},
@@ -83,14 +113,17 @@ async def txt2img(prompt: str, negative_prompt: str | None = None) -> tuple[byte
prompt_id = resp.json()["prompt_id"]
logger.info("ComfyUI queued prompt_id=%s", prompt_id)
# poll until done
for _ in range(300):
await asyncio.sleep(1)
hist = await client.get(f"{SD_BASE_URL}/history/{prompt_id}")
data = hist.json()
if prompt_id in data:
outputs = data[prompt_id]["outputs"]
# find first image output
entry = data[prompt_id]
# Log any errors from ComfyUI
if entry.get("status", {}).get("status_str") == "error":
msgs = entry.get("status", {}).get("messages", [])
logger.error("ComfyUI workflow error: %s", msgs)
outputs = entry.get("outputs", {})
for node_output in outputs.values():
if "images" in node_output:
img_info = node_output["images"][0]
@@ -100,12 +133,13 @@ async def txt2img(prompt: str, negative_prompt: str | None = None) -> tuple[byte
)
img_resp.raise_for_status()
image_bytes = img_resp.content
IMAGES_DIR.mkdir(parents=True, exist_ok=True)
filename = f"{uuid.uuid4().hex}.png"
(IMAGES_DIR / filename).write_bytes(image_bytes)
logger.info("ComfyUI done → saved %s", filename)
return image_bytes, f"images/{filename}"
logger.error("ComfyUI no image output. status=%s outputs_keys=%s",
entry.get("status"), list(outputs.keys()))
break
raise RuntimeError("ComfyUI generation timed out or produced no output")