Repository: ar-jan/llm-venice

Issue:
Add image upscaling

Discussion:

Solution patch:
diff --git a/README.md b/README.md
index f83fb10..e79946b 100644
--- a/README.md
+++ b/README.md
@@ -93,6 +93,17 @@ You can check the available parameters for a model by filtering the model list w
 
 `llm models list --query diffusion --options`
 
+### Image upscaling
+
+You can upscale existing images.
+The following saves the returned image as `image_upscaled.png` in the same directory as the original file:
+
+`llm venice upscale /path/to/image.jpg`.
+
+By default existing upscaled images are not overwritten; timestamped filenames are used instead.
+
+See `llm venice upscale --help` for the `--scale`, `--output` location, and `--overwrite` options.
+
 ### Venice commands
 
 List the available Venice commands with:
diff --git a/llm_venice.py b/llm_venice.py
index 03a981f..05ecfa3 100644
--- a/llm_venice.py
+++ b/llm_venice.py
@@ -2,6 +2,7 @@
 import datetime
 import json
 import os
+import pathlib
 from typing import Literal, Optional, Union
 
 import click
@@ -212,6 +213,68 @@ def execute(self, prompt, stream, response, conversation=None):
             raise ValueError(f"Failed to write image file: {e}")
 
 
+def image_upscale(image_path, scale, output_path=None, overwrite_files=False):
+    """
+    Upscale an image using Venice AI.
+
+    Example usage:
+        llm venice upscale image.jpg --scale 4
+    """
+    key = llm.get_key("", "venice", "LLM_VENICE_KEY")
+    if not key:
+        raise click.ClickException("No key found for Venice")
+
+    with open(image_path, "rb") as img_file:
+        image_data = img_file.read()
+
+    url = "https://api.venice.ai/api/v1/image/upscale"
+    headers = {"Authorization": f"Bearer {key}", "Accept-Encoding": "gzip"}
+
+    # Create multipart form data
+    files = {
+        "image": (pathlib.Path(image_path).name, image_data),
+    }
+
+    data = {"scale": scale}
+
+    r = httpx.post(url, headers=headers, files=files, data=data, timeout=120)
+
+    try:
+        r.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        raise ValueError(f"API request failed: {e.response.text}")
+
+    image_bytes = r.content
+
+    # Handle output path logic
+    input_path = pathlib.Path(image_path)
+    # The upscaled image is always PNG
+    default_filename = f"{input_path.stem}_upscaled.png"
+
+    if output_path is None:
+        # No output path specified, save next to input
+        output_path = input_path.parent / default_filename
+    else:
+        output_path = pathlib.Path(output_path)
+        if output_path.is_dir():
+            # If output_path is a directory, save there with default filename
+            output_path = output_path / default_filename
+
+    # Handle existing files by adding timestamp
+    if output_path.exists() and not overwrite_files:
+        stem = output_path.stem
+        suffix = output_path.suffix
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        new_filename = f"{stem}_{timestamp}{suffix}"
+        output_path = output_path.parent / new_filename
+
+    try:
+        output_path.write_bytes(image_bytes)
+        click.echo(f"Upscaled image saved to {output_path}")
+    except Exception as e:
+        raise ValueError(f"Failed to write image file: {e}")
+
+
 def refresh_models():
     "Refresh the list of models from the Venice API"
     key = llm.get_key("", "venice", "LLM_VENICE_KEY")
@@ -512,6 +575,32 @@ def new_chat(ctx, no_venice_system_prompt, web_search, character, **kwargs):
         ):
             new_chat.params.append(param)
 
+    @venice.command(name="upscale")
+    @click.argument(
+        "image_path", type=click.Path(exists=True, dir_okay=False, readable=True)
+    )
+    @click.option(
+        "--scale",
+        type=click.Choice(["2", "4"]),
+        default="2",
+        help="Scale factor (2 or 4)",
+    )
+    @click.option(
+        "--output",
+        "-o",
+        type=click.Path(file_okay=True, dir_okay=True, writable=True),
+        help="Output path (file or directory)",
+    )
+    @click.option(
+        "--overwrite",
+        is_flag=True,
+        default=False,
+        help="Overwrite existing files",
+    )
+    def upscale(image_path, scale, output, overwrite):
+        """Upscale an image using Venice API"""
+        image_upscale(image_path, scale, output, overwrite)
+
 
 @llm.hookimpl
 def register_models(register):
diff --git a/pyproject.toml b/pyproject.toml
index d706b64..8207064 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,6 +25,7 @@ venice = "llm_venice"
 test = [
     "jsonschema",
     "pytest",
+    "pytest-httpx",
     "sqlite_utils",
 ]
 


Test patch:
diff --git a/tests/test_image_upscale.py b/tests/test_image_upscale.py
new file mode 100644
index 0000000..bcb1983
--- /dev/null
+++ b/tests/test_image_upscale.py
@@ -0,0 +1,115 @@
+from unittest.mock import patch, mock_open
+
+import click
+from click.testing import CliRunner
+from llm.cli import cli
+from llm_venice import image_upscale
+import pytest
+
+
+@pytest.fixture
+def mock_image_file():
+    return b"fake image data"
+
+
+@pytest.fixture
+def mocked_responses(httpx_mock, mock_image_file):
+    """Set up mocked HTTP responses for the upscale API endpoint"""
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.venice.ai/api/v1/image/upscale",
+        content=b"upscaled image data",
+        headers={"Content-Type": "image/jpeg"},
+    )
+    return httpx_mock
+
+
+def test_upscale_function(mocked_responses, mock_image_file, tmp_path):
+    """Test the image_upscale function directly"""
+    # Create a mock image file
+    test_img_path = tmp_path / "test.jpg"
+    with open(test_img_path, "wb") as f:
+        f.write(mock_image_file)
+
+    # Test the function
+    with patch("llm.get_key", return_value="fake-key"):
+        image_upscale(str(test_img_path), 2)
+
+    # Verify the API was called correctly
+    requests = mocked_responses.get_requests()
+    assert len(requests) == 1
+    request = requests[0]
+
+    # Check the request headers
+    assert request.headers["Authorization"] == "Bearer fake-key"
+
+    # Check multipart form data
+    assert "multipart/form-data" in request.headers["Content-Type"]
+
+    # Check if the output file was created with the correct name
+    expected_output_path = tmp_path / "test_upscaled.png"
+    assert expected_output_path.exists()
+
+    # Verify the content was written
+    with open(expected_output_path, "rb") as f:
+        assert f.read() == b"upscaled image data"
+
+
+def test_upscale_command(mocked_responses, mock_image_file, tmp_path):
+    """Test the CLI command for upscaling"""
+    # Create a mock image file
+    test_img_path = tmp_path / "test.jpg"
+    with open(test_img_path, "wb") as f:
+        f.write(mock_image_file)
+
+    # Run the CLI command
+    runner = CliRunner()
+    with patch("llm.get_key", return_value="fake-key"):
+        result = runner.invoke(
+            cli, ["venice", "upscale", str(test_img_path), "--scale", "4"]
+        )
+
+    # Verify the command completed successfully
+    assert result.exit_code == 0
+
+    # Verify the output message
+    assert f"Upscaled image saved to {tmp_path}/test_upscaled.png" in result.output
+
+    # Check the request was made with the correct scale factor
+    requests = mocked_responses.get_requests()
+    assert len(requests) == 1
+
+    # Check that scale=4 was sent in the request
+    request_body = requests[0].read()
+    assert b'name="scale"' in request_body
+    assert b"4" in request_body
+
+
+def test_upscale_error_handling(httpx_mock):
+    """Test error handling in the upscale function"""
+    # Mock an error response from the API
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.venice.ai/api/v1/image/upscale",
+        status_code=400,
+        json={"error": "Invalid request"},
+    )
+
+    # Create a temporary test file
+    with patch("builtins.open", mock_open(read_data=b"fake image data")):
+        with patch("llm.get_key", return_value="fake-key"):
+            with pytest.raises(ValueError) as excinfo:
+                image_upscale("test.jpg", 2)
+
+            # Verify the error message includes the API response
+            assert "API request failed" in str(excinfo.value)
+
+
+def test_upscale_missing_api_key():
+    """Test behavior when API key is missing"""
+    # Mock get_key to return None to simulate missing API key
+    with patch("llm.get_key", return_value=None):
+        with pytest.raises(click.ClickException) as excinfo:
+            image_upscale("test.jpg", 2)
+
+        assert "No key found for Venice" in str(excinfo.value)