Add new options to generate method and CLI

Author: msabramo

firefly/cli.py

@@ -101,6 +101,15 @@ def generate(
         envvar=["FIREFLY_CLIENT_SECRET"],
     ),
     prompt: str = typer.Option(..., help="Text prompt for image generation"),
+    num_variations: int = typer.Option(None, help="Number of images to generate (numVariations)"),
+    style: str = typer.Option(None, help="Style object as JSON string (presets, imageReference, strength, etc.)"),
+    structure: str = typer.Option(None, help="Structure object as JSON string (strength, imageReference, etc.)"),
+    prompt_biasing_locale_code: str = typer.Option(None, help="Locale code for prompt biasing (promptBiasingLocaleCode)"),
+    negative_prompt: str = typer.Option(None, help="Negative prompt to avoid certain content"),
+    seed: int = typer.Option(None, help="Seed for deterministic output"),
+    aspect_ratio: str = typer.Option(None, help="Aspect ratio, e.g. '1:1', '16:9'"),
+    output_format: str = typer.Option(None, help="Output format, e.g. 'jpeg', 'png'"),
+    content_class: str = typer.Option(None, help="Content class: 'photo' or 'art'"),
     download: bool = typer.Option(False, help="Download the generated image to a file (filename is taken from the image URL)"),
     show_images: bool = typer.Option(False, help="Display the image in the terminal after download."),
     use_mocks: bool = typer.Option(False, help="Mock API responses for testing without a valid client secret."),
@@ -114,11 +123,39 @@ def generate(
         raise typer.BadParameter("client_id must be provided as an option or via the FIREFLY_CLIENT_ID environment variable.")
     if not client_secret:
         raise typer.BadParameter("client_secret must be provided as an option or via the FIREFLY_CLIENT_SECRET environment variable.")
-    with with_maybe_use_mocks(use_mocks):
-        _generate(client_id, client_secret, prompt, download, show_images, format, verbose)
-
-
-def _generate(client_id, client_secret, prompt, download, show_images, format, verbose):
+    # Parse JSON for style/structure if provided
+    style_obj = None
+    if style:
+        try:
+            style_obj = json.loads(style)
+        except Exception as e:
+            raise typer.BadParameter(f"Invalid JSON for --style: {e}")
+    structure_obj = None
+    if structure:
+        try:
+            structure_obj = json.loads(structure)
+        except Exception as e:
+            raise typer.BadParameter(f"Invalid JSON for --structure: {e}")
+    try:
+        with with_maybe_use_mocks(use_mocks):
+            _generate(
+                client_id, client_secret, prompt, download, show_images, format, verbose,
+                num_variations=num_variations,
+                style=style_obj,
+                structure=structure_obj,
+                prompt_biasing_locale_code=prompt_biasing_locale_code,
+                negative_prompt=negative_prompt,
+                seed=seed,
+                aspect_ratio=aspect_ratio,
+                output_format=output_format,
+                content_class=content_class,
+            )
+    except ValueError as e:
+        typer.secho(str(e), fg=typer.colors.RED, err=True)
+        raise typer.Exit(code=-1)
+
+
+def _generate(client_id, client_secret, prompt, download, show_images, format, verbose, **kwargs):
     client = FireflyClient(client_id=client_id, client_secret=client_secret)
     image_api_url = client.BASE_URL
     if verbose:
@@ -140,7 +177,7 @@ def capture_request(*args, **kwargs):
     _requests.request = capture_request
 
     try:
-        response = client.generate_image(prompt=prompt)
+        response = client.generate_image(prompt=prompt, **kwargs)
     finally:
         _requests.request = orig_request
 

firefly/client.py

@@ -63,13 +63,35 @@ def _request(
         except Exception as e:
             raise FireflyAPIError(f"Request failed: {e}")
 
-    def generate_image(self, prompt: str, **kwargs) -> FireflyImageResponse:
+    def generate_image(
+        self,
+        prompt: str,
+        num_variations: Optional[int] = None,
+        style: Optional[dict] = None,
+        structure: Optional[dict] = None,
+        prompt_biasing_locale_code: Optional[str] = None,
+        negative_prompt: Optional[str] = None,
+        seed: Optional[int] = None,
+        aspect_ratio: Optional[str] = None,
+        output_format: Optional[str] = None,
+        content_class: Optional[str] = None,  # 'photo' or 'art'
+        **kwargs
+    ) -> FireflyImageResponse:
         """
         Generate an image from a text prompt using Adobe Firefly.
 
         Args:
             prompt (str): The text prompt for image generation.
-            **kwargs: Additional parameters for the API (e.g., style, aspect ratio).
+            num_variations (int, optional): Number of images to generate (was `n`).
+            style (dict, optional): Style object, e.g. {"presets": [...], "imageReference": {...}, "strength": ...}.
+            structure (dict, optional): Structure reference object, e.g. {"strength": ..., "imageReference": {...}}.
+            prompt_biasing_locale_code (str, optional): Locale code for prompt biasing (was `locale`).
+            negative_prompt (str, optional): Negative prompt to avoid certain content.
+            seed (int, optional): Seed for deterministic output.
+            aspect_ratio (str, optional): Aspect ratio, e.g. "1:1", "16:9".
+            output_format (str, optional): Output format, e.g. "jpeg", "png".
+            content_class (str, optional): Content class, either 'photo' or 'art'.
+            **kwargs: Additional parameters for the API.
 
         Returns:
             FireflyImageResponse: The response object containing all fields from the API response.
@@ -78,6 +100,26 @@ def generate_image(self, prompt: str, **kwargs) -> FireflyImageResponse:
             FireflyAPIError, FireflyAuthError
         """
         data = {"prompt": prompt}
+        if num_variations is not None:
+            data["numVariations"] = num_variations
+        if style is not None:
+            data["style"] = style
+        if structure is not None:
+            data["structure"] = structure
+        if prompt_biasing_locale_code is not None:
+            data["promptBiasingLocaleCode"] = prompt_biasing_locale_code
+        if negative_prompt is not None:
+            data["negativePrompt"] = negative_prompt
+        if seed is not None:
+            data["seed"] = seed
+        if aspect_ratio is not None:
+            data["aspectRatio"] = aspect_ratio
+        if output_format is not None:
+            data["outputFormat"] = output_format
+        if content_class is not None:
+            if content_class not in ("photo", "art"):
+                raise ValueError("content_class must be either 'photo' or 'art'")
+            data["contentClass"] = content_class
         data.update(kwargs)
         resp = self._request(method="POST", url=self.BASE_URL, json=data)
         try:

tests/test_cli.py

@@ -137,3 +137,85 @@ def test_generate_verbose(monkeypatch):
     # Should include verbose status messages
     assert "Doing request to" in result.output
     assert "Received HTTP" in result.output
+
+def test_generate_with_all_new_options(monkeypatch):
+    # Valid content_class: photo
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--num-variations", "2",
+            "--style", '{"presets": ["bw"], "strength": 50}',
+            "--structure", '{"strength": 80, "imageReference": {"source": {"uploadId": "abc123"}}}',
+            "--prompt-biasing-locale-code", "en-US",
+            "--negative-prompt", "no text",
+            "--seed", "42",
+            "--aspect-ratio", "16:9",
+            "--output-format", "jpeg",
+            "--content-class", "photo",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == 0
+    assert "Generated image URL:" in result.output
+    # Valid content_class: art
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--content-class", "art",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == 0
+    assert "Generated image URL:" in result.output
+    # Invalid content_class
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--content-class", "invalid",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code != 0
+    assert "content_class must be either 'photo' or 'art'" in result.output
+
+def test_generate_invalid_json_style(monkeypatch):
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--style", "not-a-json",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == 2
+    assert "Invalid JSON for --style" in result.output
+
+def test_generate_invalid_json_structure(monkeypatch):
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--structure", "not-a-json",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == 2
+    assert "Invalid JSON for --structure" in result.output

tests/test_firefly_client.py

@@ -158,3 +158,97 @@ def test_generate_image_value_error(client, mock_valid_ims_access_token_response
     with mock.patch("requests.request", side_effect=ValueError("bad value")):
         with pytest.raises(FireflyAPIError):
             client.generate_image(prompt="trigger value error")
+
+
+@responses.activate
+def test_generate_image_with_all_new_parameters(client, mock_valid_ims_access_token_response):
+    responses.add(
+        responses.POST,
+        IMAGE_URL,
+        json={
+            "size": {"width": 1024, "height": 768},
+            "outputs": [
+                {"seed": 123, "image": {"url": "https://example.com/image.png"}}
+            ],
+            "contentClass": "photo",
+        },
+        status=200,
+    )
+    style = {"presets": ["bw"], "strength": 50}
+    structure = {"strength": 80, "imageReference": {"source": {"uploadId": "abc123"}}}
+    response = client.generate_image(
+        prompt="test prompt",
+        num_variations=2,
+        style=style,
+        structure=structure,
+        prompt_biasing_locale_code="en-US",
+        negative_prompt="no text",
+        seed=42,
+        aspect_ratio="16:9",
+        output_format="jpeg",
+        extra_param="extra_value"
+    )
+    # Check the outgoing request body
+    image_call = responses.calls[1]
+    body = json.loads(image_call.request.body.decode())
+    assert body["prompt"] == "test prompt"
+    assert body["numVariations"] == 2
+    assert body["style"] == style
+    assert body["structure"] == structure
+    assert body["promptBiasingLocaleCode"] == "en-US"
+    assert body["negativePrompt"] == "no text"
+    assert body["seed"] == 42
+    assert body["aspectRatio"] == "16:9"
+    assert body["outputFormat"] == "jpeg"
+    assert body["extra_param"] == "extra_value"
+    # Check the response is parsed correctly
+    assert response.size.width == 1024
+    assert response.size.height == 768
+    assert response.contentClass == "photo"
+    assert len(response.outputs) == 1
+    assert response.outputs[0].seed == 123
+    assert response.outputs[0].image.url == "https://example.com/image.png"
+
+
+@responses.activate
+def test_generate_image_content_class(client, mock_valid_ims_access_token_response):
+    responses.add(
+        responses.POST,
+        IMAGE_URL,
+        json={
+            "size": {"width": 512, "height": 512},
+            "outputs": [
+                {"seed": 1, "image": {"url": "https://example.com/img.png"}}
+            ],
+            "contentClass": "photo",
+        },
+        status=200,
+    )
+    # Valid value: 'photo'
+    response = client.generate_image(prompt="test", content_class="photo")
+    image_call = responses.calls[1]
+    body = json.loads(image_call.request.body.decode())
+    assert body["contentClass"] == "photo"
+    assert response.contentClass == "photo"
+    # Valid value: 'art'
+    responses.reset()
+    responses.add(
+        responses.POST,
+        IMAGE_URL,
+        json={
+            "size": {"width": 512, "height": 512},
+            "outputs": [
+                {"seed": 2, "image": {"url": "https://example.com/img2.png"}}
+            ],
+            "contentClass": "art",
+        },
+        status=200,
+    )
+    response = client.generate_image(prompt="test", content_class="art")
+    image_call = responses.calls[0]
+    body = json.loads(image_call.request.body.decode())
+    assert body["contentClass"] == "art"
+    assert response.contentClass == "art"
+    # Invalid value
+    with pytest.raises(ValueError):
+        client.generate_image(prompt="test", content_class="invalid")