Add docstrings for _models/image_generation (#115)

Author: Mandzhi

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,40 @@
 
 @dataclass(frozen=True)
 class ImageMessage:
+    """
+    This class represents a message for using in image generation models with optional weight field.
+    """
+    #: the text content of the message for using in image generation models
     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
 
-
+#: type alias for different types of messages that can be processed by image generation models
 ImageMessageType = Union[ImageMessage, ImageMessageDict, AnyMessage, str]
+#: type alias for input types accepted by the `messages_to_proto` function
 ImageMessageInputType = Union[ImageMessageType, Iterable[ImageMessageType]]
 
 
 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 of the one, concrete model. This model encapsulates the URI and configuration."""
     _config_type = ImageGenerationModelConfig
     _result_type = ImageGenerationModelResult
     _operation_type: type[OperationTypeT]
@@ -41,6 +43,16 @@ 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.
+            Read more on what MIME types exist in `the documentation <https://yandex.cloud/docs/foundation-models/image-generation/api-ref/ImageGenerationAsync/generate>`_.
+        """
         return super().configure(
             seed=seed,
             width_ratio=width_ratio,
@@ -83,6 +95,7 @@ async def _run_deferred(
             )
 
 
+@doc_from(BaseImageGenerationModel)
 class AsyncImageGenerationModel(BaseImageGenerationModel[AsyncOperation[ImageGenerationModelResult]]):
     _operation_type = AsyncOperation[ImageGenerationModelResult]
 
@@ -92,6 +105,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, or the maximum time to wait for the request to complete in seconds.
+            Defaults to 60 seconds.
+        """
         return await self._run_deferred(
             messages=messages,
             timeout=timeout
@@ -102,14 +122,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, or the maximum time to wait for the request to complete 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 +150,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,7 +15,10 @@
 
 @dataclass(frozen=True, repr=False)
 class ImageGenerationModelResult(BaseResult):
+    """This class represents the result of an image generation model inference."""
+    #: the generated image in bytes
     image_bytes: bytes
+    #: the version of the model used for generation
     model_version: str
 
     @classmethod
@@ -27,6 +30,7 @@ def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> Self:  # pylint: d
         )
 
     def _repr_jpeg_(self) -> bytes | None:
+        """:meta public:"""
         # 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