Merge f466f109bc into d9c9d84298

litellm/constants.py

@@ -1419,42 +1419,21 @@ WANDB_MODELS: Final[set] = set(
 modelscope_models: Final[set] = set(
     [
         # Qwen series models
-        "Qwen/Qwen3-0.6B",
-        "Qwen/Qwen3-1.7B",
-        "Qwen/Qwen3-4B",
-        "Qwen/Qwen3-8B",
-        "Qwen/Qwen3-14B",
-        "Qwen/Qwen3-30B-A3B",
-        "Qwen/Qwen3-32B",
-        "Qwen/Qwen3-235B-A22B",
-        "Qwen/Qwen3-235B-A22B-Instruct-2507",
-        "Qwen/Qwen3-235B-A22B-Thinking-2507",
-        "Qwen/Qwen3-30B-A3B-Thinking-2507",
-        "Qwen/Qwen3-Coder-30B-A3B-Instruct",
-        "Qwen/Qwen3-Coder-480B-A35B-Instruct",
-        "Qwen/Qwen3-Next-80B-A3B-Instruct",
-        "Qwen/Qwen3-Next-80B-A3B-Thinking",
-        "Qwen/Qwen3-VL-235B-A22B-Instruct",
-        "Qwen/Qwen3-VL-8B-Instruct",
-        "Qwen/Qwen3-VL-8B-Thinking",
         "Qwen/Qwen3.5-122B-A10B",
         "Qwen/Qwen3.5-27B",
         "Qwen/Qwen3.5-35B-A3B",
         "Qwen/Qwen3.5-397B-A17B",
-        "Qwen/QwQ-32B",
-        "Qwen/QwQ-32B-Preview",
-        "Qwen/QVQ-72B-Preview",
+        "Qwen/Qwen3.8-27B",
+        "Qwen/Qwen3.8-Flash-Next",
         "Qwen/Qwen-Image-Edit",
+        "Qwen/Qwen-Image-Edit-2511",
+        "Qwen/Qwen-Image-2512",
         # DeepSeek series models
-        "deepseek-ai/DeepSeek-R1-0528",
-        "deepseek-ai/DeepSeek-R1-Distill-Llama-70B",
-        "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B",
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
-        "deepseek-ai/DeepSeek-V3.2",
-        "deepseek-ai/DeepSeek-V4-Flash",
+        "deepseek-ai/DeepSeek-V4-Flash-0731",
+        "deepseek-ai/DeepSeek-V4-Pro",
+        "deepseek-ai/DeepSeek-V4-Pro-0813",
+        "deepseek-ai/DeepSeek-V4.1-Flash",
+        "deepseek-ai/DeepSeek-V4.1-Flash-1017",
     ]
 )
 

litellm/images/main.py

@@ -38,6 +38,7 @@ from openai.types.audio.transcription_create_params import FileTypes
 # BFL handlers
 from litellm.llms.black_forest_labs.image_edit.handler import bfl_image_edit
 from litellm.llms.black_forest_labs.image_generation.handler import bfl_image_generation
+from litellm.llms.modelscope.image_generation.handler import modelscope_image_generation
 from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config
 from litellm.main import (
     azure_chat_completions,
@@ -499,6 +500,23 @@ def image_generation(
                 client=client,
                 aimg_generation=aimg_generation,
             )
+        elif custom_llm_provider == "modelscope":
+            if api_base is not None and api_key is None:
+                raise ValueError("api_key must be provided when api_base is overridden for modelscope")
+            litellm_params_dict["api_key"] = api_key
+            litellm_params_dict["api_base"] = api_base
+            return modelscope_image_generation.image_generation(
+                model=model,
+                prompt=prompt,
+                model_response=model_response,
+                optional_params=optional_params,  # pyright: ignore[reportArgumentType]  # get_optional_params_image_gen returns dict[str, object]; values are JSON
+                litellm_params=litellm_params_dict,  # pyright: ignore[reportUnknownArgumentType]  # kwargs plumbing is untyped upstream
+                logging_obj=litellm_logging_obj,
+                timeout=timeout,
+                extra_headers=extra_headers,  # pyright: ignore[reportUnknownArgumentType]  # kwargs plumbing is untyped upstream
+                client=client,  # pyright: ignore[reportUnknownArgumentType]  # kwargs plumbing is untyped upstream
+                aimg_generation=aimg_generation,  # pyright: ignore[reportUnknownArgumentType]  # kwargs plumbing is untyped upstream
+            )
         elif custom_llm_provider == "azure_ai":
             from litellm.llms.azure_ai.common_utils import (
                 AzureFoundryModelInfo,

litellm/llms/modelscope/common_utils.py

@@ -0,0 +1,24 @@
+"""
+ModelScope Common Utilities
+
+Shared constants and error handling for ModelScope API integration.
+"""
+
+from typing import Final
+
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+
+class ModelScopeError(BaseLLMException):
+    """Exception class for ModelScope API errors."""
+
+
+DEFAULT_POLLING_INTERVAL: Final = 2.0
+DEFAULT_MAX_POLLING_TIME: Final = 300
+
+ASYNC_MODE_HEADER: Final = "X-ModelScope-Async-Mode"
+TASK_TYPE_HEADER: Final = "X-ModelScope-Task-Type"
+IMAGE_GENERATION_TASK_TYPE: Final = "image_generation"
+
+TASK_STATUS_SUCCEED: Final = "SUCCEED"
+TASK_STATUS_FAILED: Final = "FAILED"

litellm/llms/modelscope/image_generation/handler.py

@@ -0,0 +1,415 @@
+"""
+ModelScope Image Generation Handler
+
+ModelScope image generation only supports async mode: the submit call returns
+a task_id, then the caller polls GET /v1/tasks/{task_id} until task_status is
+SUCCEED (image URLs in output_images) or FAILED.
+
+API Reference: https://modelscope.cn/docs/model-service/API-Inference/intro
+"""
+
+import asyncio
+import json
+import time
+from collections.abc import Coroutine, Mapping
+from typing import Final
+
+import httpx
+from pydantic import ConfigDict, JsonValue, TypeAdapter
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
+    HTTPHandler,
+    get_async_httpx_client,  # pyright: ignore[reportUnknownVariableType]  # litellm client factory is untyped
+    get_httpx_client,  # pyright: ignore[reportUnknownVariableType]  # litellm client factory is untyped
+)
+from litellm.llms.modelscope.common_utils import (
+    DEFAULT_MAX_POLLING_TIME,
+    DEFAULT_POLLING_INTERVAL,
+    TASK_STATUS_FAILED,
+    TASK_STATUS_SUCCEED,
+    ModelScopeError,
+)
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import ImageResponse
+
+from .transformation import ModelScopeImageGenerationConfig
+
+_JSON_RESPONSE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(hide_input_in_errors=True))
+
+
+def _str_or_none(value: JsonValue) -> str | None:
+    return value if isinstance(value, str) else None
+
+
+def _extract_params(
+    litellm_params: GenericLiteLLMParams | dict[str, JsonValue],  # mutable-ok: litellm passes a mutable dict
+) -> tuple[str | None, str | None, dict[str, JsonValue]]:  # mutable-ok: returns the caller's dict unchanged
+    if isinstance(litellm_params, dict):
+        return (
+            _str_or_none(litellm_params.get("api_key")),
+            _str_or_none(litellm_params.get("api_base")),
+            litellm_params,
+        )
+    return litellm_params.api_key, litellm_params.api_base, dict(litellm_params)
+
+
+class ModelScopeImageGeneration:
+    """ModelScope image generation handler."""
+
+    def __init__(self) -> None:
+        self.config: Final = ModelScopeImageGenerationConfig()
+
+    def image_generation(
+        self,
+        model: str,
+        prompt: str,
+        model_response: ImageResponse,
+        optional_params: dict[str, JsonValue],  # mutable-ok: litellm override signature
+        litellm_params: GenericLiteLLMParams | dict[str, JsonValue],  # mutable-ok: litellm override signature
+        logging_obj: LiteLLMLoggingObj,
+        timeout: float | httpx.Timeout | None,
+        extra_headers: Mapping[str, str] | None = None,
+        client: HTTPHandler | AsyncHTTPHandler | None = None,
+        aimg_generation: bool = False,
+    ) -> ImageResponse | Coroutine[object, object, ImageResponse]:
+        if aimg_generation:
+            return self.async_image_generation(
+                model=model,
+                prompt=prompt,
+                model_response=model_response,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                logging_obj=logging_obj,
+                timeout=timeout,
+                extra_headers=extra_headers,
+                client=client if isinstance(client, AsyncHTTPHandler) else None,
+            )
+
+        api_key, api_base, litellm_params_dict = _extract_params(litellm_params)
+
+        sync_client: Final = client if isinstance(client, HTTPHandler) else get_httpx_client()  # pyright: ignore[reportUnknownVariableType]  # litellm client factory is untyped
+
+        headers: Final = self.config.validate_environment(
+            api_key=api_key,
+            headers={},
+            model=model,
+            messages=[],
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+        )
+        if extra_headers:
+            headers.update(extra_headers)
+
+        complete_url: Final = self.config.get_complete_url(
+            api_base=api_base,
+            api_key=api_key,
+            model=model,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+        )
+
+        data: Final = self.config.transform_image_generation_request(
+            model=model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+            headers=headers,
+        )
+
+        logging_obj.pre_call(  # pyright: ignore[reportUnknownMemberType]  # litellm logging API is untyped
+            input=prompt,
+            api_key="",
+            additional_args={
+                "complete_input_dict": data,
+                "api_base": complete_url,
+                "headers": headers,
+            },
+        )
+
+        try:
+            response: Final = sync_client.post(  # pyright: ignore[reportUnknownMemberType]  # untyped client post
+                url=complete_url,
+                headers=headers,
+                json=data,
+                timeout=timeout,
+            )
+        except httpx.HTTPError as e:
+            raise ModelScopeError(
+                status_code=500,
+                message=f"Request failed: {e}",
+            )
+
+        final_response: Final = self._poll_for_result_sync(
+            initial_response=response,  # pyright: ignore[reportArgumentType]  # post() returns Response | None; None raises HTTPError above
+            api_base=api_base,
+            headers=headers,
+            sync_client=sync_client,
+            timeout=timeout,
+        )
+
+        return self.config.transform_image_generation_response(
+            model=model,
+            raw_response=final_response,
+            model_response=model_response,
+            logging_obj=logging_obj,
+            request_data=data,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+            encoding=None,
+        )
+
+    async def async_image_generation(
+        self,
+        model: str,
+        prompt: str,
+        model_response: ImageResponse,
+        optional_params: dict[str, JsonValue],  # mutable-ok: litellm override signature
+        litellm_params: GenericLiteLLMParams | dict[str, JsonValue],  # mutable-ok: litellm override signature
+        logging_obj: LiteLLMLoggingObj,
+        timeout: float | httpx.Timeout | None,
+        extra_headers: Mapping[str, str] | None = None,
+        client: AsyncHTTPHandler | None = None,
+    ) -> ImageResponse:
+        api_key, api_base, litellm_params_dict = _extract_params(litellm_params)
+
+        async_client: Final = client or get_async_httpx_client(  # pyright: ignore[reportUnknownVariableType]  # litellm client factory is untyped
+            llm_provider=litellm.LlmProviders.MODELSCOPE,
+        )
+
+        headers: Final = self.config.validate_environment(
+            api_key=api_key,
+            headers={},
+            model=model,
+            messages=[],
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+        )
+        if extra_headers:
+            headers.update(extra_headers)
+
+        complete_url: Final = self.config.get_complete_url(
+            api_base=api_base,
+            api_key=api_key,
+            model=model,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+        )
+
+        data: Final = self.config.transform_image_generation_request(
+            model=model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+            headers=headers,
+        )
+
+        logging_obj.pre_call(  # pyright: ignore[reportUnknownMemberType]  # litellm logging API is untyped
+            input=prompt,
+            api_key="",
+            additional_args={
+                "complete_input_dict": data,
+                "api_base": complete_url,
+                "headers": headers,
+            },
+        )
+
+        try:
+            response: Final = await async_client.post(  # pyright: ignore[reportUnknownMemberType]  # untyped client post
+                url=complete_url,
+                headers=headers,
+                json=data,
+                timeout=timeout,
+            )
+        except httpx.HTTPError as e:
+            raise ModelScopeError(
+                status_code=500,
+                message=f"Request failed: {e}",
+            )
+
+        final_response: Final = await self._poll_for_result_async(
+            initial_response=response,  # pyright: ignore[reportArgumentType]  # post() returns Response | None; None raises HTTPError above
+            api_base=api_base,
+            headers=headers,
+            async_client=async_client,
+            timeout=timeout,
+        )
+
+        return self.config.transform_image_generation_response(
+            model=model,
+            raw_response=final_response,
+            model_response=model_response,
+            logging_obj=logging_obj,
+            request_data=data,
+            optional_params=optional_params,
+            litellm_params=litellm_params_dict,
+            encoding=None,
+        )
+
+    def _poll_for_result_sync(
+        self,
+        initial_response: httpx.Response,
+        api_base: str | None,
+        headers: dict[str, str],  # mutable-ok: caller passes a mutable dict
+        sync_client: HTTPHandler,
+        max_wait: float = DEFAULT_MAX_POLLING_TIME,
+        interval: float = DEFAULT_POLLING_INTERVAL,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> httpx.Response:
+        if initial_response.status_code >= 400:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"ModelScope submit failed: {initial_response.text}",
+            )
+
+        try:
+            response_data: Final = _JSON_RESPONSE.validate_python(initial_response.json())
+        except json.JSONDecodeError as e:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"Error parsing submit response: {e}",
+            )
+
+        if "errors" in response_data:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"ModelScope error: {response_data['errors']}",
+            )
+
+        task_id: Final[str | None] = _str_or_none(response_data.get("task_id"))
+        if not task_id:
+            raise ModelScopeError(
+                status_code=500,
+                message="No task_id in ModelScope submit response",
+            )
+
+        polling_url: Final = self.config.get_task_status_url(api_base, task_id)
+        polling_headers: Final = self.config.get_polling_headers(headers)
+
+        start_time: Final = time.time()
+        verbose_logger.debug("ModelScope starting sync polling at %s", polling_url)
+
+        while time.time() - start_time < max_wait:
+            response = sync_client.get(  # pyright: ignore[reportUnknownMemberType]  # untyped client get
+                url=polling_url,
+                headers=polling_headers,
+                timeout=timeout,
+            )
+
+            if response.status_code >= 400:
+                raise ModelScopeError(
+                    status_code=response.status_code,
+                    message=f"Polling failed: {response.text}",
+                )
+
+            try:
+                data = _JSON_RESPONSE.validate_python(response.json())
+            except json.JSONDecodeError as e:
+                raise ModelScopeError(
+                    status_code=response.status_code,
+                    message=f"Error parsing poll response: {e}",
+                )
+            status = data.get("task_status")
+            verbose_logger.debug("ModelScope poll status: %s", status)
+
+            if status == TASK_STATUS_SUCCEED:
+                return response
+            elif status == TASK_STATUS_FAILED:
+                raise ModelScopeError(
+                    status_code=400,
+                    message="ModelScope image generation task failed",
+                )
+
+            time.sleep(interval)
+
+        raise ModelScopeError(
+            status_code=408,
+            message=f"Polling timed out after {max_wait} seconds",
+        )
+
+    async def _poll_for_result_async(
+        self,
+        initial_response: httpx.Response,
+        api_base: str | None,
+        headers: dict[str, str],  # mutable-ok: caller passes a mutable dict
+        async_client: AsyncHTTPHandler,
+        max_wait: float = DEFAULT_MAX_POLLING_TIME,
+        interval: float = DEFAULT_POLLING_INTERVAL,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> httpx.Response:
+        if initial_response.status_code >= 400:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"ModelScope submit failed: {initial_response.text}",
+            )
+
+        try:
+            response_data: Final = _JSON_RESPONSE.validate_python(initial_response.json())
+        except json.JSONDecodeError as e:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"Error parsing submit response: {e}",
+            )
+
+        if "errors" in response_data:
+            raise ModelScopeError(
+                status_code=initial_response.status_code,
+                message=f"ModelScope error: {response_data['errors']}",
+            )
+
+        task_id: Final[str | None] = _str_or_none(response_data.get("task_id"))
+        if not task_id:
+            raise ModelScopeError(
+                status_code=500,
+                message="No task_id in ModelScope submit response",
+            )
+
+        polling_url: Final = self.config.get_task_status_url(api_base, task_id)
+        polling_headers: Final = self.config.get_polling_headers(headers)
+
+        start_time: Final = time.time()
+        verbose_logger.debug("ModelScope starting async polling at %s", polling_url)
+
+        while time.time() - start_time < max_wait:
+            response = await async_client.get(  # pyright: ignore[reportUnknownMemberType]  # untyped client get
+                url=polling_url,
+                headers=polling_headers,  # pyright: ignore[reportArgumentType]  # untyped client headers param
+                timeout=timeout,
+            )
+
+            if response.status_code >= 400:
+                raise ModelScopeError(
+                    status_code=response.status_code,
+                    message=f"Polling failed: {response.text}",
+                )
+
+            try:
+                data = _JSON_RESPONSE.validate_python(response.json())
+            except json.JSONDecodeError as e:
+                raise ModelScopeError(
+                    status_code=response.status_code,
+                    message=f"Error parsing poll response: {e}",
+                )
+            status = data.get("task_status")
+            verbose_logger.debug("ModelScope poll status: %s", status)
+
+            if status == TASK_STATUS_SUCCEED:
+                return response
+            elif status == TASK_STATUS_FAILED:
+                raise ModelScopeError(
+                    status_code=400,
+                    message="ModelScope image generation task failed",
+                )
+
+            await asyncio.sleep(interval)
+
+        raise ModelScopeError(
+            status_code=408,
+            message=f"Polling timed out after {max_wait} seconds",
+        )
+
+
+modelscope_image_generation: Final = ModelScopeImageGeneration()

litellm/llms/modelscope/image_generation/transformation.py

@@ -10,7 +10,7 @@ from collections.abc import Iterable, Mapping
 from typing import TYPE_CHECKING, Final
 
 import httpx
-from pydantic import ConfigDict, TypeAdapter
+from pydantic import ConfigDict, JsonValue, TypeAdapter
 from typing_extensions import override
 
 from litellm.litellm_core_utils.url_utils import join_api_base
@@ -18,6 +18,14 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
+from litellm.llms.modelscope.common_utils import (
+    ASYNC_MODE_HEADER,
+    IMAGE_GENERATION_TASK_TYPE,
+    TASK_STATUS_FAILED,
+    TASK_STATUS_SUCCEED,
+    TASK_TYPE_HEADER,
+    ModelScopeError,
+)
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -35,79 +43,71 @@ else:
 _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
 _DECODED_JSON: Final = TypeAdapter(object)
 _JSON_OBJECTS: Final = TypeAdapter(Iterable[Mapping[str, object]], config=ConfigDict(hide_input_in_errors=True))
+_JSON_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def _as_json_dict(value: JsonValue) -> dict[str, JsonValue]:  # mutable-ok: request body fields
+    return value if isinstance(value, dict) else {}
 
 
 class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
-    """
-    Configuration for ModelScope image generation.
-
-    Supports text-to-image models like:
-    - Qwen/Qwen-Image-Edit
-    - And other ModelScope-hosted image generation models
-    """
+    """Configuration for ModelScope image generation and editing models."""
 
     DEFAULT_BASE_URL: str = "https://api-inference.modelscope.cn/v1"
 
+    @override
     def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
-        """
-        Return list of OpenAI params supported by ModelScope.
-
-        ModelScope supports standard OpenAI image generation parameters.
-        """
-        return [
-            "n",  # Number of images to generate
-            "size",  # Size of the generated images
-            "response_format",  # url or b64_json
-            "user",  # User identifier
-        ]
+        # Only size is honored (verified via real API: n is silently ignored,
+        # response_format/user are not honored by ModelScope).
+        # Provider-specific fields (negative_prompt, seed, steps, guidance,
+        # loras, image_url) pass through via extra_body, not this list.
+        return ["size"]
 
+    @override
     def map_openai_params(
         self,
-        non_default_params: dict,
-        optional_params: dict,
+        non_default_params: dict[str, JsonValue],
+        optional_params: dict[str, JsonValue],
         model: str,
         drop_params: bool,
-    ) -> dict:
-        """
-        Map OpenAI parameters to ModelScope parameters.
-
-        ModelScope uses the same parameter names as OpenAI.
-        """
+    ) -> dict[str, JsonValue]:
         supported_params: Final = self.get_supported_openai_params(model)
         if drop_params:
             non_default_params = {k: v for k, v in non_default_params.items() if k in supported_params}
         optional_params.update(non_default_params)
         return optional_params
 
+    def _get_base_url(self, api_base: str | None) -> str:
+        base_url: Final[str] = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
+        return base_url.rstrip("/")
+
     @override
     def get_complete_url(
         self,
         api_base: str | None,
         api_key: str | None,
         model: str,
-        optional_params: dict,
-        litellm_params: dict,
+        optional_params: dict[str, JsonValue],
+        litellm_params: dict[str, JsonValue],
         stream: bool | None = None,
     ) -> str:
-        """
-        Get the complete URL for the ModelScope image generation API request.
-        """
-        base_url: Final = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL
+        return join_api_base(self._get_base_url(api_base), "images/generations")
 
-        # Return the images endpoint
-        return join_api_base(base_url, "images/generations")
+    def get_task_status_url(self, api_base: str | None, task_id: str) -> str:
+        """Build the URL used to poll an async image generation task."""
+        return join_api_base(self._get_base_url(api_base), f"tasks/{task_id}")
 
     @override
     def validate_environment(
         self,
-        headers: dict,
+        headers: dict[str, str],
         model: str,
         messages: list[AllMessageValues],
-        optional_params: dict,
-        litellm_params: dict,
+        optional_params: dict[str, JsonValue],
+        litellm_params: dict[str, JsonValue],
         api_key: str | None = None,
         api_base: str | None = None,
-    ) -> dict:
+    ) -> dict[str, str]:
         """
         Validate environment and set up headers for ModelScope.
         """
@@ -121,33 +121,40 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
         default_headers: Final = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {final_api_key}",
+            # ModelScope image-gen only supports async mode: the submit call
+            # returns a task_id and the caller must poll GET /v1/tasks/{task_id}.
+            ASYNC_MODE_HEADER: "true",
         }
 
         headers = {**headers, **default_headers}
         return headers
 
+    def get_polling_headers(self, headers: Mapping[str, str]) -> Mapping[str, str]:
+        return {
+            "Authorization": headers.get("Authorization", ""),
+            TASK_TYPE_HEADER: IMAGE_GENERATION_TASK_TYPE,
+        }
+
+    @override
     def transform_image_generation_request(
         self,
         model: str,
         prompt: str,
-        optional_params: dict,
-        litellm_params: dict,
-        headers: dict,
-    ) -> dict:
-        """
-        Transform OpenAI-style request to ModelScope request format.
+        optional_params: dict[str, JsonValue],
+        litellm_params: dict[str, JsonValue],
+        headers: dict[str, str],
+    ) -> dict[str, JsonValue]:
+        request_data: Final[dict[str, JsonValue]] = {"model": model, "prompt": prompt}  # mutable-ok: request body
 
-        ModelScope uses the same format as OpenAI for image generation.
-        """
-        # Build the request body (same as OpenAI)
-        request_data: Final[dict] = {
-            "model": model,
-            "prompt": prompt,
-        }
-
-        # Add optional params
-        for key, value in optional_params.items():
-            if key.startswith("_"):
+        # litellm wraps provider fields (e.g. image_url) in extra_body; flatten
+        # them into the body. Skip extra_headers/extra_query (litellm control
+        # params swept in alongside, may carry secrets) and model/prompt (set
+        # above from the routed values; a client-controlled extra_body must not
+        # override the authorized model or prompt).
+        extra_body_raw: Final = optional_params.get("extra_body") or {}
+        extra_body: Final[dict[str, JsonValue]] = _as_json_dict(extra_body_raw)
+        for key, value in {**optional_params, **extra_body}.items():
+            if key in ("extra_body", "extra_headers", "extra_query", "model", "prompt") or key.startswith("_"):
                 continue
             request_data[key] = value
 
@@ -160,21 +167,21 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
         raw_response: httpx.Response,
         model_response: ImageResponse,
         logging_obj: LiteLLMLoggingObj,
-        request_data: dict,
-        optional_params: dict,
-        litellm_params: dict,
+        request_data: dict[str, JsonValue],
+        optional_params: dict[str, JsonValue],
+        litellm_params: dict[str, JsonValue],
         encoding: object,
         api_key: str | None = None,
         json_mode: bool | None = None,
     ) -> ImageResponse:
         """
-        Transform ModelScope response to OpenAI-compatible ImageResponse.
+        Transform a completed ModelScope task response into an ImageResponse.
 
-        ModelScope returns the same format as OpenAI:
-        {"created": timestamp, "data": [{"url": "..."}]}
+        The polled response looks like:
+        {"task_status": "SUCCEED", "output_images": ["https://..."], ...}
         """
         try:
-            response_data: Final = _DECODED_JSON.validate_python(raw_response.json())
+            response_data: Final = _JSON_RESPONSE.validate_python(raw_response.json())
         except Exception as e:
             raise self.get_error_class(
                 error_message=f"Error parsing ModelScope response: {e}",
@@ -182,66 +189,59 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig):
                 headers=raw_response.headers,
             )
 
-        # Check for errors in response
-        response_object: Final = _JSON_OBJECT.validate_python(response_data)
-        if "error" in response_object:
-            error: Final = _JSON_OBJECT.validate_python(response_object["error"])
-            error_msg: Final = error.get("message", str(error))
+        if "errors" in response_data:
+            errors: Final = response_data["errors"]
+            error_msg: Final = errors.get("message", str(errors)) if isinstance(errors, dict) else str(errors)
             raise self.get_error_class(
                 error_message=f"ModelScope error: {error_msg}",
                 status_code=raw_response.status_code,
                 headers=raw_response.headers,
             )
 
-        # Extract images from response
-        data_list: Final = _JSON_OBJECTS.validate_python(response_object.get("data", []))
-        if not model_response.data:
-            model_response.data = []
-
-        for item in data_list:
-            image_obj = ImageObject(
-                url=item.get("url"),
-                b64_json=item.get("b64_json"),
-                revised_prompt=item.get("revised_prompt"),
+        task_status: Final = response_data.get("task_status")
+        if task_status == TASK_STATUS_FAILED:
+            raise self.get_error_class(
+                error_message="ModelScope image generation task failed",
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+        if task_status != TASK_STATUS_SUCCEED:
+            raise self.get_error_class(
+                error_message=(f"ModelScope task did not succeed: status={task_status}"),
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+
+        raw_output: Final = response_data.get("output_images")
+        output_images: Final[tuple[str, ...]] = (
+            tuple(url for url in raw_output if isinstance(url, str)) if isinstance(raw_output, list) else ()
+        )
+
+        if not model_response.data:
+            model_response.data = []  # rebind-ok: litellm response contract expects the passed-in object populated
+
+        for image_url in output_images:
+            model_response.data.append(ImageObject(url=image_url))
+
+        if not model_response.data:
+            raise self.get_error_class(
+                error_message="ModelScope task SUCCEED but no output_images",
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
             )
-            model_response.data.append(image_obj)
 
         return model_response
 
+    @override
     def get_error_class(
         self,
         error_message: str,
         status_code: int,
-        headers: dict | httpx.Headers,
+        headers: dict[str, str] | httpx.Headers,  # mutable-ok: litellm override signature
     ) -> BaseLLMException:
-        """Return the appropriate error class for ModelScope."""
-        from litellm.exceptions import (
-            AuthenticationError,
-            BadRequestError,
-            InternalServerError,
+        """Return the ModelScope error class, preserving the real status code."""
+        return ModelScopeError(
+            status_code=status_code,
+            message=error_message,
+            headers=headers,
         )
-
-        if status_code == 400:
-            return BadRequestError(
-                message=error_message,
-                model="",
-                llm_provider="modelscope",
-            )
-        elif status_code == 401:
-            return AuthenticationError(
-                message=error_message,
-                model="",
-                llm_provider="modelscope",
-            )
-        elif status_code >= 500:
-            return InternalServerError(
-                message=error_message,
-                model="",
-                llm_provider="modelscope",
-            )
-        else:
-            return BadRequestError(
-                message=error_message,
-                model="",
-                llm_provider="modelscope",
-            )