|
5 | 5 | import os
|
6 | 6 | from typing import List
|
7 | 7 |
|
| 8 | +import requests |
| 9 | + |
8 | 10 | from servicestack.clients import UploadFile
|
9 | 11 | from .config import *
|
10 |
| -from .aiserver_dtos import SpeechToText, GenerationResponse |
| 12 | +from .aiserver_dtos import SpeechToText, GenerationResponse, ImageToImage |
11 | 13 |
|
12 | 14 |
|
13 | 15 | def create_aiserver_client():
|
@@ -73,6 +75,49 @@ def test_can_speech_to_text(self):
|
73 | 75 | print("Text only:", text_only)
|
74 | 76 |
|
75 | 77 |
|
| 78 | + def test_image_to_image(self): |
| 79 | + """Test image to image functionality with file upload""" |
| 80 | + request = ImageToImage() |
| 81 | + request.positive_prompt = "A beautiful landscape painting" |
| 82 | + request.negative_prompt = "A pixelated image" |
| 83 | + |
| 84 | + request.model = "sdxl-lightning" |
| 85 | + |
| 86 | + # Open the test image file in binary read mode |
| 87 | + with open("tests/files/test_image.png", "rb") as image_file: |
| 88 | + upload = UploadFile( |
| 89 | + field_name="image", |
| 90 | + file_name="test_image.png", |
| 91 | + content_type="image/png", |
| 92 | + stream=image_file |
| 93 | + ) |
| 94 | + |
| 95 | + # Send request with file |
| 96 | + response: GenerationResponse = self.client.post_files_with_request( |
| 97 | + request_uri="/api/ImageToImage", |
| 98 | + request=request, |
| 99 | + files=upload |
| 100 | + ) |
| 101 | + |
| 102 | + # Verify response structure |
| 103 | + self.assertIsNotNone(response) |
| 104 | + self.assertTrue(hasattr(response, 'outputs')) |
| 105 | + self.assertIsInstance(response.outputs, List) |
| 106 | + self.assertEqual(len(response.outputs), 1) |
| 107 | + |
| 108 | + # Get image output |
| 109 | + image_output_url = response.outputs[0].url |
| 110 | + |
| 111 | + # Download the image output |
| 112 | + image_output = requests.get(image_output_url).content |
| 113 | + |
| 114 | + # Basic validation of output |
| 115 | + self.assertIsNotNone(image_output) |
| 116 | + |
| 117 | + # Save image to file |
| 118 | + with open("tests/files/image_output.webp", "wb") as output_file: |
| 119 | + output_file.write(image_output) |
| 120 | + |
76 | 121 | if __name__ == '__main__':
|
77 | 122 | # When running this file directly, automatically set RUN_AI_TESTS
|
78 | 123 | os.environ['RUN_AI_TESTS'] = '1'
|
|
0 commit comments