9 image generation from text prompts

.github/workflows/main.yml

@@ -33,7 +33,7 @@ jobs:
         pip install -r requirements.txt
     - name: Run tests with pytest
       run: |
-        pytest tests
+        pytest -m "not apitest"
     - name: Run lintering with ruff
       run: |
         # stop the build if there are Python syntax errors or undefined names

requirements.txt

@@ -1,7 +1,15 @@
+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
-exceptiongroup==1.1.3
-iniconfig==2.0.0
-packaging==23.1
-pluggy==1.3.0
-pytest==7.4.2
-tomli==2.0.1
+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

src/image_generation/image_generator.py

@@ -0,0 +1,64 @@
+# The class is a wrapper class for generating images from a prompt
+# The class uses the OpenAI API to generate images
+
+import openai
+import os
+
+class ImageGenerator:
+    MAX_PROMPT_LENGTH = 1000
+    MAX_IMAGE_SIZE = 1024
+
+    def get_user_prompt(self):
+        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):
+        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('Prompt must be less than 1000 characters')
+
+    def validate_size(self, width: int, height: int):
+        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 < 1:
+            raise ValueError('Width must be greater than 0')
+
+        if height < 1:
+            raise ValueError('Height must be greater than 0')
+
+        if width > self.MAX_IMAGE_SIZE:
+            raise ValueError('Width must be less than 1024')
+
+        if height > self.MAX_IMAGE_SIZE:
+            raise ValueError('Height must be less than 1024')
+
+    def generate_image(self, image_prompt: str, width: int, height: int):
+        self.validate_prompt(image_prompt)
+        self.validate_size(width, 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="1024x1024"
+        )
+        print(response)
+        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'])

tests/image_generation/image_generator_test.py

@@ -0,0 +1,126 @@
+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)
+        # Assert
+
+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)
+        # Assert
+
+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)
+        # Assert
+
+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)
+        # Assert
+
+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)
+        # Assert
+
+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
+
+@pytest.mark.apitest
+def test_image_generation_with_valid_prompt():
+    # Arrange
+
+    image_generator = ImageGenerator()
+    valid_prompt = "A picture of a cat"
+    # Act
+    result = image_generator.generate_image(valid_prompt, 1, 1)
+    # Assert
+    assert result is not None