Download backend/app/edit_flags.py from ruslanmv/HomePilot: direct link, hf CLI and curl.
- Browser
- Download file 13.4 kB
-
https://huggingface.co/spaces/ruslanmv/HomePilot/resolve/main/backend/app/edit_flags.py
- Command line
-
hf download hf://spaces/ruslanmv/HomePilot/backend/app/edit_flags.py
-
curl -L -o edit_flags.py https://huggingface.co/spaces/ruslanmv/HomePilot/resolve/main/backend/app/edit_flags.py
13.4 kB
| """ | |
| Edit flags parser for advanced edit controls. | |
| This module allows the frontend to pass edit parameters via hidden flags | |
| appended to the edit instruction, without modifying the /chat API schema. | |
| Example: | |
| User input: "Make the sky dramatic --steps 30 --cfg 6.5 --denoise 0.55" | |
| Parsed result: | |
| - clean_text: "Make the sky dramatic" | |
| - flags: EditFlags(steps=30, cfg=6.5, denoise=0.55, ...) | |
| """ | |
| import re | |
| from dataclasses import dataclass, field | |
| from typing import List, Optional, Tuple, Dict, Any | |
| # Pattern to match flags like --key value | |
| FLAG_RE = re.compile(r"(--[a-zA-Z0-9_-]+)\s+([^\s-][^\n\r]*?)(?=\s+--|$)") | |
| # Pattern to match URLs (http/https) | |
| URL_RE = re.compile(r"https?://\S+", re.I) | |
| # Pattern to match edit command keywords at the start | |
| EDIT_CMD_RE = re.compile(r"^\s*(edit|inpaint|modify)\s+", re.I) | |
| class EditFlags: | |
| """ | |
| Container for parsed edit flags. | |
| All values have sensible defaults for quality image editing. | |
| """ | |
| # Edit mode: auto, global, inpaint | |
| mode: str = "auto" | |
| # Generation parameters | |
| steps: int = 30 | |
| cfg: float = 5.5 | |
| denoise: float = 0.55 | |
| seed: int = 0 | |
| sampler_name: str = "euler" | |
| scheduler: str = "normal" | |
| # Model controls | |
| ckpt_name: Optional[str] = None | |
| controlnet_name: Optional[str] = None | |
| controlnet_strength: float = 1.0 | |
| cn_enabled: bool = False | |
| # Optional mask (URL or filename; backend will preprocess) | |
| mask_url: Optional[str] = None | |
| # LoRA adapters: list of {"id": str, "weight": float} | |
| loras: List[Dict[str, Any]] = field(default_factory=list) | |
| def to_dict(self) -> Dict[str, Any]: | |
| """Convert to dictionary for workflow variables.""" | |
| return { | |
| "mode": self.mode, | |
| "steps": self.steps, | |
| "cfg": self.cfg, | |
| "denoise": self.denoise, | |
| "seed": self.seed, | |
| "sampler_name": self.sampler_name, | |
| "scheduler": self.scheduler, | |
| "ckpt_name": self.ckpt_name, | |
| "controlnet_name": self.controlnet_name, | |
| "controlnet_strength": self.controlnet_strength, | |
| "cn_enabled": self.cn_enabled, | |
| "mask_url": self.mask_url, | |
| "loras": self.loras, | |
| } | |
| def parse_edit_flags(text: str) -> Tuple[str, EditFlags]: | |
| """ | |
| Extract flags from user text and return (clean_text, flags). | |
| Cleans the prompt by removing: | |
| - Edit command prefix (edit, inpaint, modify) | |
| - URLs (image URLs that are handled separately) | |
| - Flags like --steps 30, --cfg 6.5 | |
| Args: | |
| text: User's edit instruction potentially containing flags and URLs | |
| Returns: | |
| Tuple of (cleaned text without flags/URLs, parsed EditFlags object) | |
| """ | |
| flags = EditFlags() | |
| clean = text | |
| # First, parse all flags | |
| matches = list(FLAG_RE.finditer(text)) | |
| for m in matches: | |
| key = (m.group(1) or "").strip().lower() | |
| value = (m.group(2) or "").strip() | |
| try: | |
| if key == "--mode": | |
| flags.mode = value.lower() | |
| elif key == "--steps": | |
| flags.steps = int(float(value)) | |
| elif key == "--cfg": | |
| flags.cfg = float(value) | |
| elif key == "--denoise": | |
| flags.denoise = float(value) | |
| elif key == "--seed": | |
| flags.seed = int(float(value)) | |
| elif key == "--sampler": | |
| flags.sampler_name = value | |
| elif key == "--scheduler": | |
| flags.scheduler = value | |
| elif key == "--ckpt": | |
| flags.ckpt_name = value | |
| elif key == "--cn": | |
| flags.cn_enabled = value.lower() in ("1", "true", "yes", "on") | |
| elif key == "--cn-strength": | |
| flags.controlnet_strength = float(value) | |
| elif key == "--controlnet": | |
| flags.controlnet_name = value | |
| elif key == "--mask": | |
| flags.mask_url = value | |
| elif key == "--lora": | |
| # Format: "lora_id:weight" or just "lora_id" (default weight 0.8) | |
| for lora_token in value.split(): | |
| if ":" in lora_token: | |
| lora_id, lora_w = lora_token.rsplit(":", 1) | |
| try: | |
| w = float(lora_w) | |
| except ValueError: | |
| w = 0.8 | |
| else: | |
| lora_id = lora_token | |
| w = 0.8 | |
| lora_id = lora_id.strip() | |
| if lora_id: | |
| flags.loras.append({"id": lora_id, "weight": w, "enabled": True}) | |
| except (ValueError, TypeError): | |
| # Skip invalid flag values | |
| pass | |
| # Remove all flags from the text | |
| for m in reversed(matches): | |
| clean = clean[:m.start()] + clean[m.end():] | |
| # Remove URLs from the prompt (image URL is passed separately) | |
| clean = URL_RE.sub("", clean) | |
| # Remove edit command prefix (e.g., "edit ", "inpaint ") | |
| clean = EDIT_CMD_RE.sub("", clean) | |
| # Clean up extra whitespace | |
| clean = " ".join(clean.split()).strip() | |
| return clean, flags | |
| def infer_edit_mode(user_text: str) -> str: | |
| """ | |
| Heuristic routing: inpaint for localized edits, global for style changes. | |
| Args: | |
| user_text: Cleaned user edit instruction | |
| Returns: | |
| Inferred edit mode: "inpaint" or "global" | |
| """ | |
| t = user_text.lower() | |
| # Localized edit keywords suggest inpainting | |
| inpaint_keywords = ( | |
| "remove", "erase", "delete", "replace", "change", "add", "swap", | |
| "object", "logo", "text", "person", "background", "hat", "shirt", | |
| "hair", "face", "eye", "nose", "mouth", "hand", "arm", "leg", | |
| "building", "car", "tree", "sky", "water", "cloud" | |
| ) | |
| # Global style keywords suggest full image regeneration | |
| global_keywords = ( | |
| "cinematic", "anime", "oil painting", "watercolor", "cartoon", | |
| "night", "sunset", "dramatic lighting", "color grade", "style", | |
| "filter", "tone", "mood", "atmosphere", "aesthetic", "artistic", | |
| "vintage", "retro", "modern", "futuristic", "cyberpunk", "noir" | |
| ) | |
| # Check for localized edit keywords | |
| for kw in inpaint_keywords: | |
| if kw in t: | |
| return "inpaint" | |
| # Check for global style keywords | |
| for kw in global_keywords: | |
| if kw in t: | |
| return "global" | |
| # Default to auto (let the workflow decide) | |
| return "auto" | |
| def build_edit_workflow_vars( | |
| image_url: str, | |
| prompt: str, | |
| flags: EditFlags, | |
| negative_prompt: str = "", | |
| img_model: str = "", | |
| ) -> Dict[str, Any]: | |
| """ | |
| Build workflow variables dictionary from parsed flags. | |
| Args: | |
| image_url: URL of the source image | |
| prompt: Cleaned edit instruction | |
| flags: Parsed edit flags | |
| negative_prompt: Negative prompt (optional) | |
| img_model: User-selected image model from frontend (e.g. "dreamshaper_8.safetensors") | |
| Returns: | |
| Dictionary of variables for ComfyUI workflow | |
| """ | |
| import random | |
| # Default checkpoints based on workflow type | |
| # If the user selected a specific model in the frontend, use that as | |
| # the global default — this ensures SD1.5 LoRAs are not silently | |
| # skipped when the user has a SD1.5 checkpoint selected. | |
| ckpt_default_global = img_model or "sd_xl_base_1.0.safetensors" | |
| ckpt_default_inpaint = "sd_xl_base_1.0_inpainting_0.1.safetensors" | |
| ckpt_default_sd15_inpaint = "sd-v1-5-inpainting.ckpt" | |
| cn_default = "control_v11p_sd15_inpaint.safetensors" | |
| # Generate random seed if not explicitly set (prevents ComfyUI caching) | |
| seed = flags.seed | |
| if seed == 0 or seed == -1: | |
| seed = random.randint(1, 2147483647) | |
| vars_dict: Dict[str, Any] = { | |
| "image_path": image_url, | |
| "prompt": prompt, | |
| "negative_prompt": negative_prompt, | |
| "steps": flags.steps, | |
| "cfg": flags.cfg, | |
| "seed": seed, | |
| "denoise": flags.denoise, | |
| "sampler_name": flags.sampler_name, | |
| "scheduler": flags.scheduler, | |
| "filename_prefix": "homepilot_edit", | |
| } | |
| # Set checkpoint and mask based on whether mask is provided | |
| # IMPORTANT: Only set mask_path when there's actually a mask | |
| if flags.mask_url: | |
| # Inpaint mode - mask is provided | |
| vars_dict["mask_path"] = flags.mask_url | |
| vars_dict["ckpt_name"] = flags.ckpt_name or ckpt_default_inpaint | |
| # Optional ControlNet (recommended with SD1.5 inpaint) | |
| if flags.cn_enabled: | |
| vars_dict["ckpt_name"] = flags.ckpt_name or ckpt_default_sd15_inpaint | |
| vars_dict["controlnet_name"] = flags.controlnet_name or cn_default | |
| vars_dict["controlnet_strength"] = flags.controlnet_strength | |
| else: | |
| # Global edit (img2img) - no mask needed | |
| vars_dict["ckpt_name"] = flags.ckpt_name or ckpt_default_global | |
| # Pass LoRA list (will be consumed by comfy.py to inject LoraLoader nodes) | |
| # Auto-detect: if no --lora flags given, scan installed LoRAs for trigger word matches | |
| # Only selects LoRAs compatible with the current checkpoint architecture. | |
| if not flags.loras: | |
| try: | |
| from .models.lora_loader import scan_installed_loras, get_lora_dir, is_lora_compatible | |
| from .models.lora_registry import get_lora_by_id | |
| from .model_config import get_architecture | |
| lora_dir = get_lora_dir() | |
| if lora_dir.exists(): | |
| prompt_lower = vars_dict["prompt"].lower() | |
| ckpt_name = vars_dict.get("ckpt_name", "") | |
| ckpt_arch = get_architecture(ckpt_name) if ckpt_name else "" | |
| installed = scan_installed_loras() | |
| # Track which trigger words already matched (pick best arch variant) | |
| matched_triggers: set = set() | |
| for lora_info in installed: | |
| if not lora_info.get("healthy", True): | |
| continue | |
| lid = lora_info.get("id", "") | |
| entry = get_lora_by_id(lid) | |
| if not entry or not entry.trigger_words: | |
| continue | |
| # Check if any trigger word appears in the prompt | |
| matched = any(tw.lower() in prompt_lower for tw in entry.trigger_words if len(tw) >= 3) | |
| if not matched: | |
| continue | |
| # Skip if incompatible with current checkpoint | |
| if ckpt_arch and entry.base: | |
| if not is_lora_compatible(entry.base, ckpt_arch): | |
| print(f"[LORA] Skipping auto-detect '{lid}' (base={entry.base}) — incompatible with checkpoint '{ckpt_arch}'") | |
| continue | |
| # Avoid duplicates: if another LoRA with the same trigger words | |
| # was already added, skip (e.g. undressing_sd15 vs undressing_sdxl) | |
| tw_key = frozenset(tw.lower() for tw in entry.trigger_words) | |
| if tw_key in matched_triggers: | |
| continue | |
| matched_triggers.add(tw_key) | |
| flags.loras.append({"id": lid, "weight": 0.8, "enabled": True}) | |
| print(f"[LORA] Auto-detected '{lid}' (base={entry.base}) from trigger words in prompt") | |
| # Cap auto-detected LoRAs at 4 | |
| flags.loras = flags.loras[:4] | |
| except (ImportError, Exception) as e: | |
| print(f"[LORA] Auto-detect skipped: {e}") | |
| if flags.loras: | |
| vars_dict["_loras"] = flags.loras | |
| # Inject trigger words into prompt (industry best practice for LoRA activation) | |
| try: | |
| from .models.lora_registry import get_lora_by_id | |
| trigger_parts: list[str] = [] | |
| for lr in flags.loras: | |
| entry = get_lora_by_id(lr["id"]) | |
| if entry and entry.trigger_words: | |
| for tw in entry.trigger_words: | |
| if tw.lower() not in vars_dict["prompt"].lower(): | |
| trigger_parts.append(tw) | |
| if trigger_parts: | |
| vars_dict["prompt"] = ", ".join(trigger_parts) + ", " + vars_dict["prompt"] | |
| except ImportError: | |
| pass | |
| return vars_dict | |
| def determine_workflow(flags: EditFlags, prompt: str) -> str: | |
| """ | |
| Determine which workflow to use based on flags and prompt. | |
| IMPORTANT: Inpaint workflows REQUIRE a mask. If no mask is provided, | |
| we MUST use the standard edit (img2img) workflow, even if keywords | |
| suggest inpainting. The keyword-based mode detection is informational | |
| only - actual workflow selection depends on mask availability. | |
| Args: | |
| flags: Parsed edit flags | |
| prompt: Cleaned edit instruction | |
| Returns: | |
| Workflow name to use | |
| """ | |
| # Critical: Only use inpaint workflows when a mask is actually provided | |
| # The InpaintModelConditioning node requires a noise_mask input | |
| if flags.mask_url: | |
| if flags.cn_enabled: | |
| return "edit_inpaint_cn" | |
| return "edit_inpaint" | |
| # Without a mask, use standard img2img edit workflow | |
| # This applies regardless of detected keywords (remove, replace, etc.) | |
| return "edit" | |