first commit

2026-05-28 08:42:46 +03:00
commit e5c0df308f
38 changed files with 2753 additions and 0 deletions
@@ -0,0 +1,121 @@
+import asyncio
+import logging
+import os
+import uuid
+from pathlib import Path
+
+import httpx
+from dotenv import load_dotenv
+
+load_dotenv()
+
+logger = logging.getLogger(__name__)
+
+SD_BASE_URL = os.getenv("SD_BASE_URL", "http://127.0.0.1:8188").rstrip("/")
+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_DEFAULT_NEGATIVE = os.getenv(
+    "SD_DEFAULT_NEGATIVE",
+    "low quality, worst quality, blurry, bad anatomy, watermark, text",
+)
+IMAGES_DIR = Path(os.getenv("IMAGES_DIR", "static/images"))
+
+
+def split_prompt_and_negative(full_prompt: str) -> tuple[str, str]:
+    if "\n\nNegative prompt:" in full_prompt:
+        pos, _, neg = full_prompt.partition("\n\nNegative prompt:")
+        return pos.strip(), neg.strip()
+    return full_prompt.strip(), SD_DEFAULT_NEGATIVE
+
+
+def _build_workflow(positive: str, negative: str) -> dict:
+    """Minimal KSampler workflow for ComfyUI API."""
+    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]}},
+        "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,
+                "denoise": 1.0,
+            },
+        },
+    }
+
+
+async def check_sd() -> bool:
+    try:
+        async with httpx.AsyncClient(timeout=5) as client:
+            r = await client.get(f"{SD_BASE_URL}/system_stats")
+            return r.status_code == 200
+    except Exception:
+        return False
+
+
+async def txt2img(prompt: str, negative_prompt: str | None = None) -> tuple[bytes, str]:
+    neg = negative_prompt or SD_DEFAULT_NEGATIVE
+    workflow = _build_workflow(prompt, neg)
+    client_id = uuid.uuid4().hex
+
+    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},
+        )
+        resp.raise_for_status()
+        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
+                for node_output in outputs.values():
+                    if "images" in node_output:
+                        img_info = node_output["images"][0]
+                        img_resp = await client.get(
+                            f"{SD_BASE_URL}/view",
+                            params={"filename": img_info["filename"], "subfolder": img_info.get("subfolder", ""), "type": img_info.get("type", "output")},
+                        )
+                        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}"
+                break
+
+    raise RuntimeError("ComfyUI generation timed out or produced no output")
+
+
+async def generate_from_full_prompt(full_prompt: str) -> tuple[str | None, str | None]:
+    positive, negative = split_prompt_and_negative(full_prompt)
+    try:
+        _, rel_path = await txt2img(positive, negative)
+        return rel_path, None
+    except Exception as e:
+        logger.error("ComfyUI error: %s", e)
+        return None, str(e)