diff --git a/decart/models.py b/decart/models.py index be35664..a4b3bdf 100644 --- a/decart/models.py +++ b/decart/models.py @@ -183,21 +183,21 @@ class ImageToImageInput(DecartBaseModel): "lucy-2.1": ModelDefinition( name="lucy-2.1", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-2.1-vton": ModelDefinition( name="lucy-2.1-vton", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-restyle-2": ModelDefinition( name="lucy-restyle-2", url_path="/v1/stream", - fps=22, + fps=30, width=1280, height=704, ), @@ -205,35 +205,35 @@ class ImageToImageInput(DecartBaseModel): "lucy-latest": ModelDefinition( name="lucy-latest", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-vton-latest": ModelDefinition( name="lucy-vton-latest", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-vton-2": ModelDefinition( name="lucy-vton-2", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-2.1-vton-2": ModelDefinition( name="lucy-2.1-vton-2", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "lucy-restyle-latest": ModelDefinition( name="lucy-restyle-latest", url_path="/v1/stream", - fps=22, + fps=30, width=1280, height=704, ), @@ -241,14 +241,14 @@ class ImageToImageInput(DecartBaseModel): "lucy-vton": ModelDefinition( name="lucy-vton", url_path="/v1/stream", - fps=20, + fps=30, width=1088, height=624, ), "mirage_v2": ModelDefinition( name="mirage_v2", url_path="/v1/stream", - fps=22, + fps=30, width=1280, height=704, ), diff --git a/decart/realtime/client.py b/decart/realtime/client.py index f7bffda..6cc62d1 100644 --- a/decart/realtime/client.py +++ b/decart/realtime/client.py @@ -105,6 +105,8 @@ async def connect( ) -> "RealtimeClient": ws_url = f"{base_url}{options.model.url_path}" ws_url += f"?api_key={quote(api_key)}&model={quote(options.model.name)}" + if options.resolution is not None: + ws_url += f"&resolution={quote(options.resolution)}" config = WebRTCConfiguration( webrtc_url=ws_url, diff --git a/decart/realtime/types.py b/decart/realtime/types.py index 1f8cc4c..3313d2a 100644 --- a/decart/realtime/types.py +++ b/decart/realtime/types.py @@ -18,3 +18,4 @@ class RealtimeConnectOptions: on_remote_stream: Callable[[MediaStreamTrack], None] initial_state: Optional[ModelState] = None customize_offer: Optional[Callable] = None + resolution: Optional[Literal["720p", "1080p"]] = None diff --git a/tests/test_models.py b/tests/test_models.py index 3cfe7bb..420b404 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -7,20 +7,20 @@ def test_canonical_realtime_models() -> None: model = models.realtime("lucy-restyle-2") assert model.name == "lucy-restyle-2" - assert model.fps == 22 + assert model.fps == 30 assert model.width == 1280 assert model.height == 704 assert model.url_path == "/v1/stream" model = models.realtime("lucy-2.1") assert model.name == "lucy-2.1" - assert model.fps == 20 + assert model.fps == 30 assert model.width == 1088 assert model.height == 624 model = models.realtime("lucy-2.1-vton") assert model.name == "lucy-2.1-vton" - assert model.fps == 20 + assert model.fps == 30 assert model.width == 1088 assert model.height == 624 @@ -113,21 +113,21 @@ def test_latest_realtime_models() -> None: model = models.realtime("lucy-latest") assert model.name == "lucy-latest" assert model.url_path == "/v1/stream" - assert model.fps == 20 + assert model.fps == 30 assert model.width == 1088 assert model.height == 624 model = models.realtime("lucy-vton-latest") assert model.name == "lucy-vton-latest" assert model.url_path == "/v1/stream" - assert model.fps == 20 + assert model.fps == 30 assert model.width == 1088 assert model.height == 624 model = models.realtime("lucy-restyle-latest") assert model.name == "lucy-restyle-latest" assert model.url_path == "/v1/stream" - assert model.fps == 22 + assert model.fps == 30 assert model.width == 1280 assert model.height == 704 diff --git a/tests/test_realtime_unit.py b/tests/test_realtime_unit.py index bd8c78a..aec3b41 100644 --- a/tests/test_realtime_unit.py +++ b/tests/test_realtime_unit.py @@ -27,14 +27,14 @@ def test_realtime_models_available(): """Test that realtime models are available""" model = models.realtime("lucy-restyle-2") assert model.name == "lucy-restyle-2" - assert model.fps == 22 + assert model.fps == 30 assert model.width == 1280 assert model.height == 704 assert model.url_path == "/v1/stream" model2 = models.realtime("lucy-2.1") assert model2.name == "lucy-2.1" - assert model2.fps == 20 + assert model2.fps == 30 assert model2.width == 1088 assert model2.height == 624 assert model2.url_path == "/v1/stream" @@ -149,6 +149,63 @@ async def test_realtime_connect_accepts_custom_model_definition(): await realtime_client.disconnect() +async def _connect_and_capture_url(resolution=None) -> str: + client = DecartClient(api_key="test-key") + + with ( + patch("decart.realtime.client.WebRTCManager") as mock_manager_class, + patch("decart.realtime.client.aiohttp.ClientSession") as mock_session_cls, + ): + mock_manager = AsyncMock() + mock_manager.connect = AsyncMock(return_value=True) + mock_manager.is_connected = MagicMock(return_value=True) + mock_manager_class.return_value = mock_manager + + mock_session = MagicMock() + mock_session.closed = False + mock_session.close = AsyncMock() + mock_session_cls.return_value = mock_session + + from decart.realtime.types import RealtimeConnectOptions + + kwargs = {"resolution": resolution} if resolution is not None else {} + realtime_client = await RealtimeClient.connect( + base_url=client.realtime_base_url, + api_key=client.api_key, + local_track=MagicMock(), + options=RealtimeConnectOptions( + model=models.realtime("lucy-2.1"), + on_remote_stream=lambda t: None, + **kwargs, + ), + ) + + call_args = mock_manager_class.call_args + config = call_args[0][0] if call_args[0] else call_args[1]["configuration"] + url = config.webrtc_url + + await realtime_client.disconnect() + return url + + +@pytest.mark.asyncio +async def test_realtime_connect_omits_resolution_when_unset(): + url = await _connect_and_capture_url() + assert "resolution=" not in url + + +@pytest.mark.asyncio +async def test_realtime_connect_appends_resolution_720p(): + url = await _connect_and_capture_url("720p") + assert "&resolution=720p" in url + + +@pytest.mark.asyncio +async def test_realtime_connect_appends_resolution_1080p(): + url = await _connect_and_capture_url("1080p") + assert "&resolution=1080p" in url + + @pytest.mark.asyncio async def test_realtime_set_prompt_with_mock(): """Test set_prompt with mocked WebRTC and prompt_ack""" diff --git a/uv.lock b/uv.lock index 47915c7..b5b8a4c 100644 --- a/uv.lock +++ b/uv.lock @@ -597,7 +597,7 @@ wheels = [ [[package]] name = "decart" -version = "0.0.37" +version = "0.0.38" source = { editable = "." } dependencies = [ { name = "aiofiles" },