feat: Add avatar-live and reference image support

- Add avatar-live realtime model with set_image() method
- Add lucy-restyle-v2v video model with reference_image support
- Add AvatarOptions for configuring avatar image on connect
- Add VideoRestyleInput with mutual exclusivity validation
- Add comprehensive tests for both features

Author: YehoavR

decart/__init__.py

@@ -12,7 +12,7 @@
     QueueResultError,
     TokenCreateError,
 )
-from .models import models, ModelDefinition
+from .models import models, ModelDefinition, VideoRestyleInput
 from .types import FileInput, ModelState, Prompt
 from .queue import (
     QueueClient,
@@ -31,6 +31,7 @@
         RealtimeClient,
         RealtimeConnectOptions,
         ConnectionState,
+        AvatarOptions,
     )
 
     REALTIME_AVAILABLE = True
@@ -39,6 +40,7 @@
     RealtimeClient = None  # type: ignore
     RealtimeConnectOptions = None  # type: ignore
     ConnectionState = None  # type: ignore
+    AvatarOptions = None  # type: ignore
 
 __version__ = "0.0.1"
 
@@ -56,6 +58,7 @@
     "QueueResultError",
     "models",
     "ModelDefinition",
+    "VideoRestyleInput",
     "FileInput",
     "ModelState",
     "Prompt",
@@ -75,5 +78,6 @@
             "RealtimeClient",
             "RealtimeConnectOptions",
             "ConnectionState",
+            "AvatarOptions",
         ]
     )

decart/models.py

@@ -1,10 +1,10 @@
 from typing import Literal, Optional, List, Generic, TypeVar
-from pydantic import BaseModel, Field, ConfigDict
+from pydantic import BaseModel, Field, ConfigDict, model_validator
 from .errors import ModelNotFoundError
 from .types import FileInput, MotionTrajectoryInput
 
 
-RealTimeModels = Literal["mirage", "mirage_v2", "lucy_v2v_720p_rt"]
+RealTimeModels = Literal["mirage", "mirage_v2", "lucy_v2v_720p_rt", "avatar-live"]
 VideoModels = Literal[
     "lucy-dev-i2v",
     "lucy-fast-v2v",
@@ -13,6 +13,7 @@
     "lucy-pro-v2v",
     "lucy-pro-flf2v",
     "lucy-motion",
+    "lucy-restyle-v2v",
 ]
 ImageModels = Literal["lucy-pro-t2i", "lucy-pro-i2i"]
 Model = Literal[RealTimeModels, VideoModels, ImageModels]
@@ -95,6 +96,38 @@ class ImageToMotionVideoInput(DecartBaseModel):
     resolution: Optional[str] = None
 
 
+class VideoRestyleInput(DecartBaseModel):
+    """Input for lucy-restyle-v2v model.
+
+    Must provide either `prompt` OR `reference_image`, but not both.
+    `enhance_prompt` is only valid when using `prompt`, not `reference_image`.
+    """
+
+    prompt: Optional[str] = Field(default=None, min_length=1, max_length=1000)
+    reference_image: Optional[FileInput] = None
+    data: FileInput
+    seed: Optional[int] = None
+    resolution: Optional[str] = None
+    enhance_prompt: Optional[bool] = None
+
+    @model_validator(mode="after")
+    def validate_prompt_or_reference_image(self) -> "VideoRestyleInput":
+        has_prompt = self.prompt is not None
+        has_reference_image = self.reference_image is not None
+
+        if has_prompt == has_reference_image:
+            raise ValueError(
+                "Must provide either 'prompt' or 'reference_image', but not both"
+            )
+
+        if has_reference_image and self.enhance_prompt is not None:
+            raise ValueError(
+                "'enhance_prompt' is only valid when using 'prompt', not 'reference_image'"
+            )
+
+        return self
+
+
 class TextToImageInput(BaseModel):
     prompt: str = Field(
         ...,
@@ -144,6 +177,14 @@ class ImageToImageInput(DecartBaseModel):
             height=704,
             input_schema=BaseModel,
         ),
+        "avatar-live": ModelDefinition(
+            name="avatar-live",
+            url_path="/v1/avatar-live/stream",
+            fps=25,
+            width=1280,
+            height=720,
+            input_schema=BaseModel,
+        ),
     },
     "video": {
         "lucy-dev-i2v": ModelDefinition(
@@ -202,6 +243,14 @@ class ImageToImageInput(DecartBaseModel):
             height=704,
             input_schema=ImageToMotionVideoInput,
         ),
+        "lucy-restyle-v2v": ModelDefinition(
+            name="lucy-restyle-v2v",
+            url_path="/v1/generate/lucy-restyle-v2v",
+            fps=25,
+            width=1280,
+            height=704,
+            input_schema=VideoRestyleInput,
+        ),
     },
     "image": {
         "lucy-pro-t2i": ModelDefinition(
@@ -247,6 +296,7 @@ def video(model: VideoModels) -> VideoModelDefinition:
             - "lucy-dev-i2v" - Image-to-video (Dev quality)
             - "lucy-fast-v2v" - Video-to-video (Fast quality)
             - "lucy-motion" - Image-to-motion-video
+            - "lucy-restyle-v2v" - Video-to-video with prompt or reference image
         """
         try:
             return _MODELS["video"][model]  # type: ignore[return-value]

decart/queue/request.py

@@ -24,7 +24,7 @@ async def submit_job(
 
     for key, value in inputs.items():
         if value is not None:
-            if key in ("data", "start", "end"):
+            if key in ("data", "start", "end", "reference_image"):
                 content, content_type = await file_input_to_bytes(value, session)
                 form_data.add_field(key, content, content_type=content_type)
             else:

decart/realtime/__init__.py

@@ -1,8 +1,9 @@
 from .client import RealtimeClient
-from .types import RealtimeConnectOptions, ConnectionState
+from .types import RealtimeConnectOptions, ConnectionState, AvatarOptions
 
 __all__ = [
     "RealtimeClient",
     "RealtimeConnectOptions",
     "ConnectionState",
+    "AvatarOptions",
 ]

decart/realtime/client.py

@@ -1,21 +1,33 @@
 from typing import Callable, Optional
 import asyncio
+import base64
 import logging
 import uuid
+import aiohttp
 from aiortc import MediaStreamTrack
 
 from .webrtc_manager import WebRTCManager, WebRTCConfiguration
-from .messages import PromptMessage
+from .messages import PromptMessage, SetAvatarImageMessage
 from .types import ConnectionState, RealtimeConnectOptions
+from ..types import FileInput
 from ..errors import DecartSDKError, InvalidInputError, WebRTCError
+from ..process.request import file_input_to_bytes
 
 logger = logging.getLogger(__name__)
 
 
 class RealtimeClient:
-    def __init__(self, manager: WebRTCManager, session_id: str):
+    def __init__(
+        self,
+        manager: WebRTCManager,
+        session_id: str,
+        http_session: Optional[aiohttp.ClientSession] = None,
+        is_avatar_live: bool = False,
+    ):
         self._manager = manager
         self.session_id = session_id
+        self._http_session = http_session
+        self._is_avatar_live = is_avatar_live
         self._connection_callbacks: list[Callable[[ConnectionState], None]] = []
         self._error_callbacks: list[Callable[[DecartSDKError], None]] = []
 
@@ -24,14 +36,16 @@ async def connect(
         cls,
         base_url: str,
         api_key: str,
-        local_track: MediaStreamTrack,
+        local_track: Optional[MediaStreamTrack],
         options: RealtimeConnectOptions,
         integration: Optional[str] = None,
     ) -> "RealtimeClient":
         session_id = str(uuid.uuid4())
         ws_url = f"{base_url}{options.model.url_path}"
         ws_url += f"?api_key={api_key}&model={options.model.name}"
 
+        is_avatar_live = options.model.name == "avatar-live"
+
         config = WebRTCConfiguration(
             webrtc_url=ws_url,
             api_key=api_key,
@@ -43,16 +57,33 @@ async def connect(
             initial_state=options.initial_state,
             customize_offer=options.customize_offer,
             integration=integration,
+            is_avatar_live=is_avatar_live,
         )
 
+        # Create HTTP session for file conversions
+        http_session = aiohttp.ClientSession()
+
         manager = WebRTCManager(config)
-        client = cls(manager=manager, session_id=session_id)
+        client = cls(
+            manager=manager,
+            session_id=session_id,
+            http_session=http_session,
+            is_avatar_live=is_avatar_live,
+        )
 
         config.on_connection_state_change = client._emit_connection_change
         config.on_error = lambda error: client._emit_error(WebRTCError(str(error), cause=error))
 
         try:
-            await manager.connect(local_track)
+            # For avatar-live, convert and send avatar image before WebRTC connection
+            avatar_image_base64: Optional[str] = None
+            if is_avatar_live and options.avatar:
+                image_bytes, _ = await file_input_to_bytes(
+                    options.avatar.avatar_image, http_session
+                )
+                avatar_image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+
+            await manager.connect(local_track, avatar_image_base64=avatar_image_base64)
 
             if options.initial_state:
                 if options.initial_state.prompt:
@@ -61,6 +92,7 @@ async def connect(
                         enrich=options.initial_state.prompt.enrich,
                     )
         except Exception as e:
+            await http_session.close()
             raise WebRTCError(str(e), cause=e)
 
         return client
@@ -100,6 +132,47 @@ async def set_prompt(self, prompt: str, enrich: bool = True) -> None:
         finally:
             self._manager.unregister_prompt_wait(prompt)
 
+    async def set_image(self, image: FileInput) -> None:
+        """Set or update the avatar image.
+
+        Only available for avatar-live model.
+
+        Args:
+            image: The image to set. Can be bytes, Path, URL string, or file-like object.
+
+        Raises:
+            InvalidInputError: If not using avatar-live model or image is invalid.
+            DecartSDKError: If the server fails to acknowledge the image.
+        """
+        if not self._is_avatar_live:
+            raise InvalidInputError("set_image() is only available for avatar-live model")
+
+        if not self._http_session:
+            raise InvalidInputError("HTTP session not available")
+
+        # Convert image to base64
+        image_bytes, _ = await file_input_to_bytes(image, self._http_session)
+        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+
+        event, result = self._manager.register_image_set_wait()
+
+        try:
+            await self._manager.send_message(
+                SetAvatarImageMessage(type="set_image", image_data=image_base64)
+            )
+
+            try:
+                await asyncio.wait_for(event.wait(), timeout=15.0)
+            except asyncio.TimeoutError:
+                raise DecartSDKError("Image set acknowledgment timed out")
+
+            if not result["success"]:
+                raise DecartSDKError(
+                    result.get("status") or "Failed to set avatar image"
+                )
+        finally:
+            self._manager.unregister_image_set_wait()
+
     def is_connected(self) -> bool:
         return self._manager.is_connected()
 
@@ -108,6 +181,8 @@ def get_connection_state(self) -> ConnectionState:
 
     async def disconnect(self) -> None:
         await self._manager.cleanup()
+        if self._http_session and not self._http_session.closed:
+            await self._http_session.close()
 
     def on(self, event: str, callback: Callable) -> None:
         if event == "connection_change":

decart/realtime/messages.py

@@ -51,9 +51,22 @@ class PromptAckMessage(BaseModel):
     error: Optional[str] = None
 
 
+class ImageSetMessage(BaseModel):
+    """Acknowledgment for avatar image set from server."""
+
+    type: Literal["image_set"]
+    status: str
+
+
 # Discriminated union for incoming messages
 IncomingMessage = Annotated[
-    Union[AnswerMessage, IceCandidateMessage, SessionIdMessage, PromptAckMessage],
+    Union[
+        AnswerMessage,
+        IceCandidateMessage,
+        SessionIdMessage,
+        PromptAckMessage,
+        ImageSetMessage,
+    ],
     Field(discriminator="type"),
 ]
 
@@ -79,8 +92,17 @@ class PromptMessage(BaseModel):
     enhance_prompt: bool = True
 
 
+class SetAvatarImageMessage(BaseModel):
+    """Set avatar image message."""
+
+    type: Literal["set_image"]
+    image_data: str  # Base64-encoded image
+
+
 # Outgoing message union (no discriminator needed - we know what we're sending)
-OutgoingMessage = Union[OfferMessage, IceCandidateMessage, PromptMessage]
+OutgoingMessage = Union[
+    OfferMessage, IceCandidateMessage, PromptMessage, SetAvatarImageMessage
+]
 
 
 def parse_incoming_message(data: dict) -> IncomingMessage:

decart/realtime/types.py

@@ -1,7 +1,7 @@
 from typing import Literal, Callable, Optional
 from dataclasses import dataclass
 from ..models import ModelDefinition
-from ..types import ModelState
+from ..types import ModelState, FileInput
 
 try:
     from aiortc import MediaStreamTrack
@@ -12,9 +12,18 @@
 ConnectionState = Literal["connecting", "connected", "disconnected"]
 
 
+@dataclass
+class AvatarOptions:
+    """Options for avatar-live model."""
+
+    avatar_image: FileInput
+    """The avatar image to use. Can be bytes, Path, URL string, or file-like object."""
+
+
 @dataclass
 class RealtimeConnectOptions:
     model: ModelDefinition
     on_remote_stream: Callable[[MediaStreamTrack], None]
     initial_state: Optional[ModelState] = None
     customize_offer: Optional[Callable] = None
+    avatar: Optional[AvatarOptions] = None