style: Refactor code formatting and improve readability in nodes.py

feat: Add Ruff workflow for code linting
Refactor ModelScope Image Generator API (#5 )
2026-04-10 17:48:02 +01:00 · 2026-04-10 17:46:42 +01:00 · 2026-01-05 02:26:23 +00:00
3 changed files with 244 additions and 73 deletions
@@ -0,0 +1,25 @@
+name: Ruff
+
+on:
+  push:
+    branches: ["master"]
+  pull_request:
+
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install Ruff
+        run: pip install ruff
+
+      - name: Run Ruff check
+        run: ruff check .
@@ -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
@@ -4,22 +4,32 @@ import os
 import time
 import requests
 import importlib
+import json
+import logging

 BASE_URL = "https://api-inference.modelscope.cn/v1"

-MODELSCOPE_API_KEY = os.getenv("MODELSCOPE_API_KEY") or os.getenv("MODELSCOPE_ACCESS_TOKEN")
+MODELSCOPE_API_KEY = os.getenv("MODELSCOPE_API_KEY") or os.getenv(
+    "MODELSCOPE_ACCESS_TOKEN"
+)

-DEFAULT_SYSTEM_PROMPT = "You are a helpful and harmless assistant. Answer concisely and helpfully."
+DEFAULT_SYSTEM_PROMPT = (
+    "You are a helpful and harmless assistant. Answer concisely and helpfully."
+)

 # Constants
 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:
    """Base class for ModelScope nodes with common functionality."""
-    
+
    def _resolve_key(self, api_key: Optional[str]) -> str:
        """Resolve API key from input or environment variables."""
        key = api_key or MODELSCOPE_API_KEY
@@ -31,41 +41,35 @@ class ModelScopeBase:

    def _sanitize_text(self, text: str) -> str:
        """Sanitize text for ComfyUI showText node compatibility.
-        
+
        Args:
            text: Raw text string
-            
+
        Returns:
            str: Sanitized text safe for display
        """
        if not isinstance(text, str):
            text = str(text)
-        
-        # Remove or replace problematic characters
+
        import re
-        
-        # Remove null bytes and other control characters except common ones
-        text = re.sub(r'[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]', '', text)
-        
-        # Normalize whitespace
-        text = re.sub(r'\s+', ' ', text).strip()
-        
-        # Ensure text is not empty
+
+        text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]", "", text)
+
+        text = re.sub(r"\s+", " ", text).strip()
+
        if not text:
            text = "[Empty Response]"
-        
-        # Limit length if extremely long (optional safeguard)
+
        max_length = 10000
        if len(text) > max_length:
            text = text[:max_length] + "...[truncated]"
-        
-        # Ensure proper encoding
+
        try:
-            text.encode('utf-8').decode('utf-8')
+            text.encode("utf-8").decode("utf-8")
        except UnicodeError:
            # Fallback: replace problematic characters
-            text = text.encode('utf-8', errors='replace').decode('utf-8')
-        
+            text = text.encode("utf-8", errors="replace").decode("utf-8")
+
        return text

    def _get_pil_image_module(self):
@@ -75,24 +79,23 @@ 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]
-            
+
        Raises:
            RuntimeError: If download or conversion fails
        """
        try:
            import torch
            import numpy as np
-            
-            # Download image
+
            response = requests.get(url, headers=headers, timeout=30)
            response.raise_for_status()
            
@@ -109,19 +112,18 @@ class ModelScopeBase:
            
            # Convert PIL to numpy array
            image_np = np.array(image).astype(np.float32) / 255.0
-            
-            # Convert to torch tensor with ComfyUI format [B, H, W, C]
+
            image_tensor = torch.from_numpy(image_np)[None,]
-            
+
            return image_tensor
-            
+
        except Exception as e:
            raise RuntimeError(f"Failed to download image: {e}") from e


 class ModelScopeChatBase(ModelScopeBase):
    """Base class for chat-based ModelScope nodes."""
-    
+
    def _extract_text(self, data: dict) -> str:
        """Extract text content from chat completion response.
        Tries to read choices[0].message.content (string or list) and falls back to str(data).
@@ -216,22 +218,21 @@ class ModelScopeLLM(ModelScopeChatBase):
        request_timeout: int = 30,
    ) -> Tuple[str]:
        """Execute text generation using ModelScope LLM.
-        
+
        Args:
            model_id: ModelScope LLM model identifier
            prompt: User text prompt for the model
            system_prompt: Optional system prompt for assistant behavior
            api_key: Optional API key override
            request_timeout: Request timeout in seconds
-            
+
        Returns:
            tuple: Single-element tuple with LLM response text
-            
+
        Raises:
            ValueError: If inputs are invalid
            RuntimeError: If API call or processing fails
        """
-        # Validate inputs
        if not prompt.strip():
            raise ValueError("Prompt cannot be empty")

@@ -297,13 +298,13 @@ class ModelScopeLLM(ModelScopeChatBase):

        elapsed = time.time() - start_time
        print(f"LLM inference completed in {elapsed:.1f}s")
-        
+
        return (text,)


 class ModelScopeVLM(ModelScopeChatBase):
    """ModelScope Visual Language Model node for image-text conversations.
-    
+
    Supports analyzing images with text prompts using ModelScope VLM models.
    Handles various image input formats and provides robust error handling.
    """
@@ -311,7 +312,7 @@ class ModelScopeVLM(ModelScopeChatBase):
    @classmethod
    def INPUT_TYPES(cls):
        """Define input types and validation rules for VLM processing.
-        
+
        Returns:
            dict: Input type definitions with validation parameters
        """
@@ -387,7 +388,7 @@ class ModelScopeVLM(ModelScopeChatBase):
        request_timeout: int = 120,
    ) -> Tuple[str]:
        """Execute vision-language model inference on image and text.
-        
+
        Args:
            model_id: ModelScope VLM model identifier
            image_url: URL of image to analyze
@@ -395,27 +396,25 @@ class ModelScopeVLM(ModelScopeChatBase):
            system_prompt: Optional system prompt for assistant behavior
            api_key: Optional API key override
            request_timeout: Request timeout in seconds
-            
+
        Returns:
            tuple: Single-element tuple with VLM response text
-            
+
        Raises:
            ValueError: If inputs are invalid
            RuntimeError: If API call or processing fails
        """
-        # Validate inputs
        if not image_url.strip():
            raise ValueError("Image URL cannot be empty")
        if not prompt.strip():
            raise ValueError("Prompt cannot be empty")
-            
-        # Validate image URL format
-        if not (image_url.startswith('http://') or image_url.startswith('https://')):
+
+        if not (image_url.startswith("http://") or image_url.startswith("https://")):
            raise ValueError("Image URL must be a valid HTTP/HTTPS URL")

        key = self._resolve_key(api_key)
        system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT
-        
+
        url = f"{BASE_URL}/chat/completions"
        headers = {
            "Authorization": f"Bearer {key}",
@@ -457,7 +456,6 @@ class ModelScopeVLM(ModelScopeChatBase):

        start_time = time.time()
        try:
-            # Use longer timeout for VLM inference
            resp = requests.post(
                url,
                headers=headers,
@@ -488,7 +486,7 @@ class ModelScopeVLM(ModelScopeChatBase):

        elapsed = time.time() - start_time
        print(f"VLM inference completed in {elapsed:.1f}s")
-        
+
        return (text,)


@@ -525,7 +523,7 @@ class ModelScopeImageGenerator(ModelScopeBase):
                        "tooltip": "Negative prompt to avoid unwanted elements.",
                    },
                ),
-                "num_inference_steps": (
+                "steps": (
                    "INT",
                    {
                        "default": 30,
@@ -535,7 +533,7 @@ class ModelScopeImageGenerator(ModelScopeBase):
                        "tooltip": "Number of denoising steps (higher = better quality, slower).",
                    },
                ),
-                "guidance_scale": (
+                "guidance": (
                    "FLOAT",
                    {
                        "default": 3.5,
@@ -545,6 +543,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 +572,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",
                    {
@@ -572,7 +587,7 @@ class ModelScopeImageGenerator(ModelScopeBase):
                        "tooltip": "ModelScope API key. If empty, reads from MODELSCOPE_API_KEY env var.",
                    },
                ),
-            }
+            },
        }

    RETURN_TYPES = ("IMAGE",)
@@ -588,63 +603,155 @@ 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.
-        
+
        Args:
            model_id: ModelScope model identifier
            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
-            
+
        Raises:
            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
+            raise RuntimeError(
+                f"Network error calling ModelScope Image API: {e}"
+            ) from e

        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 +764,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "ModelScopeLLM": "ModelScope LLM",
    "ModelScopeVLM": "ModelScope VLM",
    "ModelScopeImageGenerator": "ModelScope Image Generator",
-}
+}