feat: add arguments width and height

Author: remarkablemark

.github/workflows/test.yml

@@ -62,6 +62,7 @@ jobs:
       matrix:
         command:
           - diffused segmind/tiny-sd "an apple"
+          - diffused segmind/tiny-sd "an apple" --output apple.jpg --width=512 --height=512
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -89,7 +90,7 @@ jobs:
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
-          name: artifact
+          name: artifact-${{ github.run_id }}-${{ github.job }}
           path: |
             *.jpg
             *.png

README.md

@@ -58,13 +58,41 @@ diffused dreamlike-art/dreamlike-photoreal-2.0 "cinematic photo of Godzilla eati
 **Optional**: Generated image filename.
 
 ```sh
-diffused dreamlike-art/dreamlike-photoreal-2.0 "cat eating sushi" --output cat.jpg
+diffused dreamlike-art/dreamlike-photoreal-2.0 "cat eating sushi" --output=cat.jpg
 ```
 
 With short option:
 
 ```sh
-diffused dreamlike-art/dreamlike-photoreal-2.0 "cat eating sushi" -o cat.jpg
+diffused dreamlike-art/dreamlike-photoreal-2.0 "cat eating sushi" -o=cat.jpg
+```
+
+### `--width`
+
+**Optional**: Generated image width in pixels.
+
+```sh
+diffused stabilityai/stable-diffusion-xl-base-1.0 "dog in space" --width=1024
+```
+
+With short option:
+
+```sh
+diffused stabilityai/stable-diffusion-xl-base-1.0 "dog in space" -W=1024
+```
+
+### `--height`
+
+**Optional**: Generated image height in pixels.
+
+```sh
+diffused stabilityai/stable-diffusion-xl-base-1.0 "dog in space" --height=1024
+```
+
+With short option:
+
+```sh
+diffused stabilityai/stable-diffusion-xl-base-1.0 "dog in space" -H=1024
 ```
 
 ### `--version`

pyproject.toml

@@ -6,7 +6,7 @@ authors = [
 ]
 description = "Generate images with diffusion models"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 classifiers = [
     "Programming Language :: Python :: 3",
     "Operating System :: OS Independent",

src/diffused/cli.py

@@ -30,10 +30,26 @@ def main(argv: list[str] = None) -> None:
         help="output file",
     )
 
+    parser.add_argument(
+        "--width",
+        "-W",
+        type=int,
+        help="image width",
+    )
+
+    parser.add_argument(
+        "--height",
+        "-H",
+        type=int,
+        help="image height",
+    )
+
     args = parser.parse_args(argv)
 
     filename = args.output if args.output else f"{uuid1()}.png"
-    image = generate(model=args.model, prompt=args.prompt)
+    image = generate(
+        model=args.model, prompt=args.prompt, width=args.width, height=args.height
+    )
     image.save(filename)
     print(f"🤗 {filename}")
 

src/diffused/generate.py

@@ -2,16 +2,20 @@
 from PIL import Image
 
 
-def generate(model: str, prompt: str) -> Image.Image:
+def generate(
+    model: str, prompt: str, width: int | None = None, height: int | None = None
+) -> Image.Image:
     """
     Generate image with diffusion model.
 
     Args:
         model (str): Diffusion model.
         prompt (str): Text prompt.
+        width (int): Image width.
+        height (int): Image height.
 
     Returns:
         image (PIL.Image.Image): Pillow image.
     """
     pipeline = AutoPipelineForText2Image.from_pretrained(model)
-    return pipeline(prompt).images[0]
+    return pipeline(prompt=prompt, width=width, height=height).images[0]

tests/test_cli.py

@@ -14,22 +14,36 @@ def test_version(capsys: pytest.LogCaptureFixture) -> None:
     assert captured.out == __version__ + "\n"
 
 
+def test_version_short(capsys: pytest.LogCaptureFixture) -> None:
+    with pytest.raises(SystemExit):
+        main(["-v"])
+    captured = capsys.readouterr()
+    assert captured.out == __version__ + "\n"
+
+
 def test_help(capsys: pytest.LogCaptureFixture) -> None:
     with pytest.raises(SystemExit):
         main(["--help"])
     captured = capsys.readouterr()
     assert "Generate image with diffusion model" in captured.out
 
 
-def test_required(capsys: pytest.LogCaptureFixture) -> None:
+def test_help_short(capsys: pytest.LogCaptureFixture) -> None:
+    with pytest.raises(SystemExit):
+        main(["-h"])
+    captured = capsys.readouterr()
+    assert "Generate image with diffusion model" in captured.out
+
+
+def test_no_arguments(capsys: pytest.LogCaptureFixture) -> None:
     with pytest.raises(SystemExit) as exception:
         main([])
     captured = capsys.readouterr()
     assert exception.type is SystemExit
     assert "error: the following arguments are required: model, prompt" in captured.err
 
 
-def test_invalid(capsys: pytest.LogCaptureFixture) -> None:
+def test_invalid_argument(capsys: pytest.LogCaptureFixture) -> None:
     with pytest.raises(SystemExit) as exception:
         main(["model", "prompt", "--invalid"])
     captured = capsys.readouterr()
@@ -39,7 +53,7 @@ def test_invalid(capsys: pytest.LogCaptureFixture) -> None:
 
 @patch("diffusers.AutoPipelineForText2Image.from_pretrained")
 @patch("PIL.Image.Image.save")
-def test_generate(
+def test_generate_image(
     mock_from_pretrained: Mock, mock_save: Mock, capsys: pytest.LogCaptureFixture
 ) -> None:
     main(["model", "prompt"])
@@ -56,11 +70,45 @@ def test_generate(
 
 @patch("diffusers.AutoPipelineForText2Image.from_pretrained")
 @patch("PIL.Image.Image.save")
-def test_generate_output(
+def test_output(
     mock_from_pretrained: Mock, mock_save: Mock, capsys: pytest.LogCaptureFixture
 ) -> None:
     filename = "image.png"
     main(["model", "prompt", "--output", filename])
     mock_save.assert_called_once()
     captured = capsys.readouterr()
     assert captured.out == f"🤗 {filename}\n"
+
+
+@patch("diffusers.AutoPipelineForText2Image.from_pretrained")
+@patch("PIL.Image.Image.save")
+def test_output_short(
+    mock_from_pretrained: Mock, mock_save: Mock, capsys: pytest.LogCaptureFixture
+) -> None:
+    filename = "image.png"
+    main(["model", "prompt", "-o", filename])
+    mock_save.assert_called_once()
+    captured = capsys.readouterr()
+    assert captured.out == f"🤗 {filename}\n"
+
+
+@patch("diffusers.AutoPipelineForText2Image.from_pretrained")
+@patch("PIL.Image.Image.save")
+def test_width_height(
+    mock_from_pretrained: Mock, mock_save: Mock, capsys: pytest.LogCaptureFixture
+) -> None:
+    main(["model", "prompt", "--width", "1024", "--height", "1024"])
+    mock_save.assert_called_once()
+    captured = capsys.readouterr()
+    assert "🤗 " in captured.out
+
+
+@patch("diffusers.AutoPipelineForText2Image.from_pretrained")
+@patch("PIL.Image.Image.save")
+def test_width_height_short(
+    mock_from_pretrained: Mock, mock_save: Mock, capsys: pytest.LogCaptureFixture
+) -> None:
+    main(["model", "prompt", "-W", "1024", "-H", "1024"])
+    mock_save.assert_called_once()
+    captured = capsys.readouterr()
+    assert "🤗 " in captured.out

tests/test_generate.py

@@ -2,11 +2,9 @@
 
 from diffused import generate
 
-mock_image = Mock()
 
-
-def pipeline(prompt: str):
-    return {"images": [mock_image]}  # pragma: no cover
+def pipeline(prompt: str, width: int | None, height: int | None):
+    pass  # pragma: no cover
 
 
 mock_pipeline = create_autospec(pipeline)
@@ -20,4 +18,19 @@ def test_generate(mock_from_pretrained: Mock) -> None:
     image = generate(model=model, prompt=prompt)
     assert isinstance(image, Mock)
     mock_from_pretrained.assert_called_once_with(model)
-    mock_pipeline.assert_called_once_with(prompt)
+    mock_pipeline.assert_called_once_with(prompt=prompt, width=None, height=None)
+    mock_pipeline.reset_mock()
+
+
+@patch("diffusers.AutoPipelineForText2Image.from_pretrained")
+def test_generate_width_height(mock_from_pretrained: Mock) -> None:
+    mock_from_pretrained.return_value = mock_pipeline
+    model = "model/test"
+    prompt = "test prompt"
+    width = 1024
+    height = 1024
+    image = generate(model=model, prompt=prompt, width=width, height=height)
+    assert isinstance(image, Mock)
+    mock_from_pretrained.assert_called_once_with(model)
+    mock_pipeline.assert_called_once_with(prompt=prompt, width=width, height=height)
+    mock_pipeline.reset_mock()