Add docstrings for _models/image_generation

src/yandex_cloud_ml_sdk/_models/image_generation/function.py

@@ -3,18 +3,40 @@
 from typing_extensions import override
 
 from yandex_cloud_ml_sdk._types.function import BaseModelFunction, ModelTypeT
+from yandex_cloud_ml_sdk._utils.doc import doc_from
 
 from .model import AsyncImageGenerationModel, ImageGenerationModel
 
 
 class BaseImageGeneration(BaseModelFunction[ModelTypeT]):
+    """
+    A class for image generation models.
+
+    It provides the functionality to call an image generation model by constructing
+    the appropriate URI based on the provided model name and version.
+
+    Returns a model's object through which requests to the backend are made.
+
+    >>> model = sdk.models.image_generation('yandex-art')  # this is how the model is created
+    """
     @override
     def __call__(
         self,
         model_name: str,
         *,
         model_version: str = 'latest',
     ):
+        """
+        Call the image generation model with the specified name and version.
+
+        Constructs the URI for the model based on the provided model's name
+        and version. If the name contains '://', it is treated as a
+        complete URI. Otherwise, it constructs a URI using the folder ID
+        from the SDK.
+
+        :param model_name: the name of the image generation model.
+        :param model_version: the version of the model to use (default is 'latest').
+        """
         if '://' in model_name:
             uri = model_name
         else:
@@ -26,10 +48,10 @@ def __call__(
             uri=uri,
         )
 
-
+@doc_from(BaseImageGeneration)
 class ImageGeneration(BaseImageGeneration):
     _model_type = ImageGenerationModel
 
-
+@doc_from(BaseImageGeneration)
 class AsyncImageGeneration(BaseImageGeneration):
     _model_type = AsyncImageGenerationModel

src/yandex_cloud_ml_sdk/_models/image_generation/message.py

@@ -12,26 +12,41 @@
 
 @dataclass(frozen=True)
 class ImageMessage:
+    """
+    This class represents an image message with optional weight field.
+    """
+    #: the text content of the image message
     text: str
+    #: the weight associated with the message
     weight: float | None = None
 
 
 class ImageMessageDict(TypedDict):
+    """
+    The class with TypedDict representing the structure of an image message.
+    """
     text: str
     weight: NotRequired[float]
 
 
 # NB: it supports _messages.message.Message and _models.completions.message.TextMessage
 @runtime_checkable
 class AnyMessage(Protocol):
+    """
+    The class with a protocol which defines an object with a text field.
+    The protocol can be used to check if an object has a text attribute.
+    """
     text: str
 
 
 ImageMessageType = Union[ImageMessage, ImageMessageDict, AnyMessage, str]
+"""Type alias for different types of image messages that can be processed."""
 ImageMessageInputType = Union[ImageMessageType, Iterable[ImageMessageType]]
+"""Type alias for input types accepted by the `messages_to_proto` function."""
 
 
 def messages_to_proto(messages: ImageMessageInputType) -> list[ProtoMessage]:
+    """:meta private:"""
     msgs: tuple[ImageMessageType] = coerce_tuple(  # type: ignore[assignment]
         messages,
         (dict, str, ImageMessage, AnyMessage)  # type: ignore[arg-type]

src/yandex_cloud_ml_sdk/_models/image_generation/model.py

@@ -18,6 +18,7 @@
 from yandex_cloud_ml_sdk._types.misc import UNDEFINED, UndefinedOr
 from yandex_cloud_ml_sdk._types.model import ModelAsyncMixin, OperationTypeT
 from yandex_cloud_ml_sdk._types.operation import AsyncOperation, Operation
+from yandex_cloud_ml_sdk._utils.doc import doc_from
 from yandex_cloud_ml_sdk._utils.sync import run_sync
 
 from .config import ImageGenerationModelConfig
@@ -28,6 +29,7 @@
 class BaseImageGenerationModel(
     ModelAsyncMixin[ImageGenerationModelConfig, ImageGenerationModelResult, OperationTypeT],
 ):
+    """A class for image generation models."""
     _config_type = ImageGenerationModelConfig
     _result_type = ImageGenerationModelResult
     _operation_type: type[OperationTypeT]
@@ -41,6 +43,15 @@ def configure(  # type: ignore[override]
         height_ratio: UndefinedOr[int] = UNDEFINED,
         mime_type: UndefinedOr[str] = UNDEFINED,
     ) -> Self:
+        """
+        Configures the image generation model with specified parameters and
+        returns the configured instance of the model.
+
+        :param seed: a random seed for generation.
+        :param width_ratio: the width ratio for the generated image.
+        :param height_ratio: the height ratio for the generated image.
+        :param mime_type: the MIME type of the generated image.
+        """
         return super().configure(
             seed=seed,
             width_ratio=width_ratio,
@@ -83,6 +94,7 @@ async def _run_deferred(
             )
 
 
+@doc_from(BaseImageGenerationModel)
 class AsyncImageGenerationModel(BaseImageGenerationModel[AsyncOperation[ImageGenerationModelResult]]):
     _operation_type = AsyncOperation[ImageGenerationModelResult]
 
@@ -92,6 +104,13 @@ async def run_deferred(
         *,
         timeout: float = 60,
     ) -> AsyncOperation[ImageGenerationModelResult]:
+        """Executes the image generation operation asynchronously
+        and returns an operation representing the ongoing image generation process.
+
+        :param messages: the input messages for image generation.
+        :param timeout: the timeout for the operation in seconds.
+            Defaults to 60 seconds.
+        """
         return await self._run_deferred(
             messages=messages,
             timeout=timeout
@@ -102,14 +121,22 @@ async def attach_deferred(
         operation_id: str,
         timeout: float = 60,
     ) -> AsyncOperation[ImageGenerationModelResult]:
+        """Attaches to an ongoing image generation operation.
+
+        :param operation_id: the ID of the operation to attach to.
+        :param timeout: the timeout for the operation in seconds.
+            Defaults to 60 seconds.
+        """
         return await self._attach_deferred(operation_id=operation_id, timeout=timeout)
 
 
+@doc_from(BaseImageGenerationModel)
 class ImageGenerationModel(BaseImageGenerationModel[Operation[ImageGenerationModelResult]]):
     _operation_type = Operation[ImageGenerationModelResult]
     __run_deferred = run_sync(BaseImageGenerationModel[Operation[ImageGenerationModelResult]]._run_deferred)
     __attach_deferred = run_sync(BaseImageGenerationModel[Operation[ImageGenerationModelResult]]._attach_deferred)
 
+    @doc_from(AsyncImageGenerationModel.run_deferred)
     def run_deferred(
         self,
         messages: ImageMessageInputType,
@@ -122,6 +149,7 @@ def run_deferred(
             timeout=timeout
         )
 
+    @doc_from(AsyncImageGenerationModel.attach_deferred)
     def attach_deferred(self, operation_id: str, timeout: float = 60) -> Operation[ImageGenerationModelResult]:
         return cast(
             Operation[ImageGenerationModelResult],

src/yandex_cloud_ml_sdk/_models/image_generation/result.py

@@ -15,18 +15,23 @@
 
 @dataclass(frozen=True, repr=False)
 class ImageGenerationModelResult(BaseResult):
+    """This class represents the result of an image generation model."""
+    #: the generated image in bytes
     image_bytes: bytes
+    #: the version of the model used for generation
     model_version: str
 
     @classmethod
     def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> Self:  # pylint: disable=unused-argument
+        """:meta private:"""
         proto = cast(ImageGenerationResponse, proto)
         return cls(
             image_bytes=proto.image,
             model_version=proto.model_version,
         )
 
     def _repr_jpeg_(self) -> bytes | None:
+        """:meta private:"""
         # NB: currently model could return only jpeg,
         # but for future I will put this check here to
         # remember we will need to make a _repr_png_ and such
@@ -38,5 +43,6 @@ def _repr_jpeg_(self) -> bytes | None:
         return None
 
     def __repr__(self) -> str:
+        """:meta private:"""
         size = len(self.image_bytes)
         return f'{self.__class__.__name__}(model_version={self.model_version!r}, image_bytes=<{size} bytes>)'