From 7e210fe7776754951d5ebfcc2a2e4d9c2e4a2687 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Mon, 5 Jan 2026 02:26:23 +0000 Subject: [PATCH] Refactor ModelScope Image Generator API (#5) * Refactor ModelScope Image Generator to use new Async API. Updates `ModelScopeImageGenerator` in `nodes.py` to: - Use `v1/images/generations` with `X-ModelScope-Async-Mode: true`. - Implement polling logic for task completion. - Update input parameters to match API documentation (steps, guidance, seed, loras, size). - Support `loras` as string or JSON dict. * Initial plan * Address all review comments from PR review thread Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> * Add .gitignore and remove __pycache__ Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> * Move logging import to top and improve headers copy Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> * Define logger at module level and fix formatting Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> * Refactor ModelScope Image Generator API Updates `ModelScopeImageGenerator` to use the asynchronous `v1/images/generations` endpoint. - Implements task polling logic with timestamp-based timeout. - Adds fail-fast logic for 4xx client errors during polling. - Updates input parameters to match API: `steps`, `guidance`, `seed`, `loras`. - Supports `loras` as JSON dict or string ID. - Uses `logging` instead of `print`. - Ensures UTF-8 charset in headers for multi-language support. * Simplify loras parsing logic to avoid duplication Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: neverbiasu <90372299+neverbiasu@users.noreply.github.com> --- .gitignore | 39 ++++++++++++++ nodes.py | 155 +++++++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 173 insertions(+), 21 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..228d580 --- /dev/null +++ b/.gitignore @@ -0,0 +1,39 @@ +# Python cache +__pycache__/ +*.py[cod] +*$py.class +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual environments +venv/ +ENV/ +env/ + +# IDEs +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db diff --git a/nodes.py b/nodes.py index a88713b..895e998 100644 --- a/nodes.py +++ b/nodes.py @@ -4,6 +4,8 @@ import os import time import requests import importlib +import json +import logging BASE_URL = "https://api-inference.modelscope.cn/v1" @@ -15,6 +17,10 @@ DEFAULT_SYSTEM_PROMPT = "You are a helpful and harmless assistant. Answer concis PLACEHOLDER_MODEL_ID = "Model ID" PLACEHOLDER_API_KEY = "API Key" CONTENT_TYPE_JSON = "application/json" +CONTENT_TYPE_JSON_UTF8 = "application/json; charset=utf-8" + +# Logger +logger = logging.getLogger(__name__) class ModelScopeBase: @@ -75,12 +81,12 @@ class ModelScopeBase: except Exception: raise RuntimeError("Pillow is required. Install with: pip install pillow") - def _download_image_from_url(self, url: str, headers: dict) -> object: + def _download_image_from_url(self, url: str, headers: Optional[dict] = None) -> object: """Download image from URL and convert to ComfyUI IMAGE format. Args: url: Image URL to download - headers: HTTP headers for authentication + headers: HTTP headers for authentication (optional) Returns: torch.Tensor: Image tensor in ComfyUI format [B, H, W, C] @@ -525,7 +531,7 @@ class ModelScopeImageGenerator(ModelScopeBase): "tooltip": "Negative prompt to avoid unwanted elements.", }, ), - "num_inference_steps": ( + "steps": ( "INT", { "default": 30, @@ -535,7 +541,7 @@ class ModelScopeImageGenerator(ModelScopeBase): "tooltip": "Number of denoising steps (higher = better quality, slower).", }, ), - "guidance_scale": ( + "guidance": ( "FLOAT", { "default": 3.5, @@ -545,6 +551,15 @@ class ModelScopeImageGenerator(ModelScopeBase): "tooltip": "How closely to follow the prompt (higher = more adherent).", }, ), + "seed": ( + "INT", + { + "default": 1234, + "min": 0, + "max": 0xffffffffffffffff, + "tooltip": "Random seed for reproducibility.", + }, + ), "height": ( "INT", { @@ -565,6 +580,14 @@ class ModelScopeImageGenerator(ModelScopeBase): "tooltip": "Generated image width in pixels.", }, ), + "loras": ( + "STRING", + { + "default": "", + "placeholder": '{"lora_id": 0.6} or just lora_id', + "tooltip": "LoRA model configuration.", + } + ), "api_key": ( "STRING", { @@ -588,10 +611,12 @@ class ModelScopeImageGenerator(ModelScopeBase): prompt: str, api_key: Optional[str] = None, negative_prompt: Optional[str] = None, - num_inference_steps: int = 30, - guidance_scale: float = 3.5, + steps: int = 30, + guidance: float = 3.5, + seed: Optional[int] = None, height: int = 1024, width: int = 1024, + loras: Optional[str] = None, ) -> Tuple[object]: """Generate image using ModelScope API. @@ -600,10 +625,12 @@ class ModelScopeImageGenerator(ModelScopeBase): prompt: Text description of desired image api_key: Optional API key override negative_prompt: Elements to avoid in generation - num_inference_steps: Number of denoising steps - guidance_scale: Adherence to prompt + steps: Number of denoising steps + guidance: Adherence to prompt + seed: Random seed height: Image height in pixels width: Image width in pixels + loras: LoRA configuration string (ID or JSON) Returns: tuple: Single-element tuple containing image tensor @@ -612,23 +639,52 @@ class ModelScopeImageGenerator(ModelScopeBase): RuntimeError: If API call or image processing fails """ key = self._resolve_key(api_key) + # Note: AIGC API uses base_url without /v1/ suffix for the base, + # but the endpoint is /v1/images/generations. url = f"{BASE_URL}/images/generations" + headers = { "Authorization": f"Bearer {key}", - "Content-Type": CONTENT_TYPE_JSON + "Content-Type": CONTENT_TYPE_JSON_UTF8, + "X-ModelScope-Async-Mode": "true" } payload = { "model": model_id, "prompt": prompt, - "negative_prompt": negative_prompt, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, "size": f"{width}x{height}", } + if negative_prompt: + payload["negative_prompt"] = negative_prompt + if steps is not None: + payload["steps"] = steps + if guidance is not None: + payload["guidance"] = guidance + if seed is not None: + payload["seed"] = seed + + if loras: + loras = loras.strip() + if loras: + # Try to parse as JSON, otherwise use as string + try: + parsed_loras = json.loads(loras) + # Only accept JSON objects (dict); otherwise, use the original string + if isinstance(parsed_loras, dict): + payload["loras"] = parsed_loras + else: + payload["loras"] = loras + except json.JSONDecodeError: + # Fallback to string if parsing fails + payload["loras"] = loras + try: - resp = requests.post(url, headers=headers, json=payload, timeout=60) + # Use ensure_ascii=False for Chinese characters support + data_bytes = json.dumps(payload, ensure_ascii=False).encode('utf-8') + headers_with_charset = headers.copy() + headers_with_charset["Content-Type"] = "application/json; charset=utf-8" + resp = requests.post(url, headers=headers_with_charset, data=data_bytes, timeout=60) resp.raise_for_status() except requests.RequestException as e: raise RuntimeError(f"Network error calling ModelScope Image API: {e}") from e @@ -636,15 +692,72 @@ class ModelScopeImageGenerator(ModelScopeBase): try: data = resp.json() except Exception as e: - raise RuntimeError(f"Failed to parse API response: {e}") from e + raise RuntimeError(f"Failed to parse API response as JSON: {e}") from e - # Handle direct image URL response - if data.get("images"): - first_image = data["images"][0] - if isinstance(first_image, dict) and first_image.get("url"): - return (self._download_image_from_url(first_image["url"], headers),) + task_id = data.get("task_id") + if not task_id: + raise RuntimeError(f"No task_id returned from async request: {data}") - raise RuntimeError(f"Unable to locate generated image in response: {data}") + # Poll for status + task_url = f"{BASE_URL}/tasks/{task_id}" + poll_headers = { + "Authorization": f"Bearer {key}", + "Content-Type": CONTENT_TYPE_JSON, + "X-ModelScope-Task-Type": "image_generation" + } + + # Poll for up to 10 minutes (600s) + max_wait_time = 600 + start_poll_time = time.time() + + while time.time() - start_poll_time < max_wait_time: + try: + result = requests.get(task_url, headers=poll_headers, timeout=30) + result.raise_for_status() + task_data = result.json() + + status = task_data.get("task_status") + if status == "SUCCEED": + output_images = task_data.get("output_images") + if output_images and len(output_images) > 0: + image_url = output_images[0] + download_headers = {"Authorization": f"Bearer {key}"} if key else None + return (self._download_image_from_url(image_url, download_headers),) + else: + raise RuntimeError( + f"Task succeeded but no output images found (status={status}, " + f"output_count={len(output_images) if output_images is not None else 0})" + ) + + elif status == "FAILED": + error_message = task_data.get("error_message") or task_data.get("message") or "Unknown error" + error_code = task_data.get("error_code") or task_data.get("code") + details = f", code={error_code}" if error_code is not None else "" + raise RuntimeError(f"Image generation failed (status={status}{details}): {error_message}") + + elif status in ["PENDING", "RUNNING"]: + time.sleep(5) + continue + + else: + # Unknown status: fail fast with a clear error + raise RuntimeError(f"Unknown task status '{status}' received from API: {task_data}") + + except requests.RequestException as e: + # Distinguish permanent HTTP errors (4xx) from transient issues. + status_code = getattr(getattr(e, "response", None), "status_code", None) + if status_code is not None and 400 <= status_code < 500: + # Client errors are typically permanent for this task; fail fast. + raise RuntimeError( + f"Image generation polling failed with client error {status_code}: " + f"{getattr(e.response, 'text', '')}" + ) from e + + # Transient network or server error during polling: log and retry. + logger.warning(f"Polling failed: {e}. Retrying...") + time.sleep(5) + + raise RuntimeError("Image generation timed out after polling.") NODE_CLASS_MAPPINGS = { @@ -657,4 +770,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ModelScopeLLM": "ModelScope LLM", "ModelScopeVLM": "ModelScope VLM", "ModelScopeImageGenerator": "ModelScope Image Generator", -} \ No newline at end of file +}