diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 2123f56..ef0d7c9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -17,7 +17,7 @@ jobs: strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11"] + python-version: ["3.8", "3.9", "3.10"] steps: - uses: actions/checkout@v2 @@ -28,18 +28,8 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install flake8 pytest - pip install ruff pip install -r requirements.txt - name: Run tests with pytest run: | - pytest tests - - name: Run lintering with ruff - run: | - # stop the build if there are Python syntax errors or undefined names - flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics - # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide - flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics - - ruff src --fix # make ruff run for src folder files and --fix to fix automatically linting errors - ruff tests --fix + pytest -m "not apitest" + diff --git a/requirements.txt b/requirements.txt index e7c5a17..4d51ff3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,18 @@ +aiohttp==3.8.5 +aiosignal==1.3.1 +async-timeout==4.0.3 +attrs==23.1.0 +certifi==2023.7.22 +charset-normalizer==3.2.0 colorama==0.4.6 +frozenlist==1.4.0 +idna==3.4 +multidict==6.0.4 +openai==0.28.0 +requests==2.31.0 +tqdm==4.66.1 +urllib3==2.0.4 +yarl==1.9.2 exceptiongroup==1.1.3 iniconfig==2.0.0 packaging==23.1 diff --git a/src/image_generation/image_generator.py b/src/image_generation/image_generator.py new file mode 100644 index 0000000..51f7b07 --- /dev/null +++ b/src/image_generation/image_generator.py @@ -0,0 +1,63 @@ +# The class is a wrapper class for generating images from a prompt +# The class uses the OpenAI API to generate images + +import openai + +class ImageGenerator: + MAX_PROMPT_LENGTH = 1000 + SUPPORTED_SIZES = [256, 512, 1024] + + def get_user_prompt(self) -> str: + image_prompt = input('What do you want to generate an image of? ') + self._validate_prompt(image_prompt) + return image_prompt + + def _validate_prompt(self, image_prompt: str) -> None: + if not isinstance(image_prompt, str): + raise TypeError('Prompt must be a string') + + if len(image_prompt) == 0: + raise ValueError('Prompt must not be empty') + + if (len(image_prompt) > self.MAX_PROMPT_LENGTH): + raise ValueError(f'Prompt must be less than {self.MAX_PROMPT_LENGTH} characters') + + def _validate_size(self, width: int, height: int) -> None: + if not isinstance(width, int): + raise TypeError('Width must be an integer') + + if not isinstance(height, int): + raise TypeError('Height must be an integer') + + if (width != height): + raise ValueError('Width and height must be equal') + + if (width not in self.SUPPORTED_SIZES): + raise ValueError(f'Width must be {self.SUPPORTED_SIZES}') + + def generate_image(self, image_prompt: str, width: int, height: int) -> dict: + ''' + Generates an image from a prompt of a certain size + Args: + image_prompt (str): The prompt to generate the image from. Must be less than 1000 characters. + width (int): The width of the image. Must be 256, 512, or 1024 + height (int): The height of the image. Must be 256, 512, or 1024 + Returns: + A dictionary containing the image url + ''' + self._validate_prompt(image_prompt) + self._validate_size(width, height) + size = str(width) + 'x' + str(height) + openai.api_key = '' # TODO: Make sure this is the way we get environment variables + response = openai.Image.create( + prompt=image_prompt, + n=1, + size=size + ) + return response + +if __name__ == "__main__": + image_generator = ImageGenerator() + image_prompt = image_generator.get_user_prompt() + response = image_generator.generate_image(image_prompt, 512, 512) + print(response['data'][0]['url']) \ No newline at end of file diff --git a/tests/image_generation/image_generator_test.py b/tests/image_generation/image_generator_test.py new file mode 100644 index 0000000..c45a608 --- /dev/null +++ b/tests/image_generation/image_generator_test.py @@ -0,0 +1,150 @@ +import pytest + +from src.image_generation.image_generator import ImageGenerator + +def test_image_generation_with_empty_string_prompt(): + # Arrange + image_generator = ImageGenerator() + no_prompt = "" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(no_prompt, 1, 1) + +def test_image_generation_with_non_string_prompt(): + # Arrange + image_generator = ImageGenerator() + non_string_prompt = 1 + with pytest.raises(TypeError): + # Act + result = image_generator.generate_image(non_string_prompt, 1, 1) + +def test_image_generation_with_none_prompt(): + # Arrange + image_generator = ImageGenerator() + none_prompt = None + with pytest.raises(TypeError): + # Act + result = image_generator.generate_image(none_prompt, 1, 1) + + +def test_image_generation_with_too_long_prompt(): + # Arrange + image_generator = ImageGenerator() + too_long_prompt = "A" * 3001 + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(too_long_prompt, 1, 1) + +def test_image_generation_with_non_integer_width(): + # Arrange + image_generator = ImageGenerator() + non_integer_width = "1" + valid_height = 1 + valid_prompt = "A picture of a cat" + with pytest.raises(TypeError): + # Act + result = image_generator.generate_image(valid_prompt, non_integer_width, valid_height) + +def test_image_generation_with_non_integer_height(): + # Arrange + image_generator = ImageGenerator() + valid_width = 1 + non_integer_height = "1" + valid_prompt = "A picture of a cat" + with pytest.raises(TypeError): + # Act + result = image_generator.generate_image(valid_prompt, valid_width, non_integer_height) + # Assert + +def test_image_generation_with_too_small_width(): + # Arrange + + image_generator = ImageGenerator() + too_small_width = 0 + valid_height = 1 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, too_small_width, valid_height) + # Assert + +def test_image_generation_with_too_small_height(): + # Arrange + + image_generator = ImageGenerator() + valid_width = 1 + too_small_height = 0 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, valid_width, too_small_height) + # Assert + +def test_image_generation_with_too_large_width(): + # Arrange + + image_generator = ImageGenerator() + too_large_width = 1025 + valid_height = 1 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, too_large_width, valid_height) + # Assert + +def test_image_generation_with_too_large_height(): + # Arrange + + image_generator = ImageGenerator() + valid_width = 1 + too_large_height = 1025 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, valid_width, too_large_height) + # Assert + +def test_image_generation_with_width_not_equal_to_height(): + # Arrange + + image_generator = ImageGenerator() + width_not_equal_to_height = 512 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, width_not_equal_to_height, width_not_equal_to_height + 1) + # Assert + +def test_image_generation_with_invalid_width(): + # Arrange + + image_generator = ImageGenerator() + invalid_width = 300 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, invalid_width, invalid_width) + # Assert + +def test_image_generation_with_invalid_height(): + # Arrange + + image_generator = ImageGenerator() + invalid_height = 300 + valid_prompt = "A picture of a cat" + with pytest.raises(ValueError): + # Act + result = image_generator.generate_image(valid_prompt, invalid_height, invalid_height) + # Assert + +@pytest.mark.apitest +def test_image_generation_with_valid_prompt(): + # Arrange + + image_generator = ImageGenerator() + valid_size = 256 + valid_prompt = "A picture of a cat" + # Act + result = image_generator.generate_image(valid_prompt, valid_size, valid_size) + # Assert + assert result is not None \ No newline at end of file