Upscaling image utility

1
import requests
2
import json
3
import os
4
import base64
5
import time
6
import io
7
import PIL.Image
8
from typing import Optional, Tuple
9
10
11
def _process_test(url):
12
    
13
    image = PIL.Image.open("headphones1.jpeg")
14
15
    # Create a BytesIO buffer to hold the image data
16
    image_buffer = io.BytesIO()
17
    image.save(image_buffer, format='JPEG')
18
    image_bytes = image_buffer.getvalue()
19
    encoded_image = base64.b64encode(image_bytes).decode('utf-8')
20
21
    OCTOAI_TOKEN = os.environ.get("OCTOAI_TOKEN")
22
23
    payload = {
24
        "model": "real-esrgan-x4-plus",
25
        "scale": 2,
26
        "init_image": encoded_image,
27
        "output_image_encoding": "jpeg"
28
    }
29
    headers = {
30
        "Authorization": f"Bearer {OCTOAI_TOKEN}",
31
        "Content-Type": "application/json",
32
        "X-OctoAI-Queue-Dispatch": "true"
33
    }
34
35
    response = requests.post(url, headers=headers, json=payload)
36
37
    if response.status_code != 200:
38
        print(response.text)
39
    print(response.json())
40
41
    img_info = response.json()["image_b64"]
42
43
    img_bytes = base64.b64decode(img_info)
44
    img = PIL.Image.open(io.BytesIO(img_bytes))
45
46
    if img.mode == 'RGBA':
47
     img = img.convert('RGB')
48
49
    img.save("result_image.jpeg")
50
51
if __name__ == "__main__":
52
    _process_test("https://image.octoai.run/upscaling")