import numpy as np
import gradio as gr
import requests
import io
import base64
import PIL
from PIL import Image

url = "https://api.runpod.ai/v2/d4p75k87yroni1/runsync"

def convert(input_img, quality=85):
    buffer = io.BytesIO()
    input_img.save(buffer, format="JPEG", quality=quality)
    buffer.seek(0)
    img_base64 = base64.b64encode(buffer.read()).decode('utf-8')
    
    return img_base64

def send_req(input_img, compression, noise):

    if type(input_img) is not PIL.Image.Image:
        input_img = Image.fromarray(input_img, 'RGB')
    
    payload = {
      "input": {
        "image": convert(input_img),
        "mode": "1",
        "quality": str(compression),
        "noise": str(noise)
      }
    }

    headers = {
        "Authorization": "Bearer XWV1ST04C0QLWNVAUSJWI6VJMR7YDJCKJSAR6TPA",
        "content-type": "application/json"
    }
    
    response = requests.post(url, json=payload, headers=headers)

    image_data = base64.b64decode(response.json()["output"])
    image = Image.open(io.BytesIO(image_data))
    
    return image

demo = gr.Interface(send_req, [gr.Image(), gr.Slider(0, 100, label="Compression", step=1), gr.Slider(0, 100, label="Noise", step=1)], "image")
if __name__ == "__main__":
    demo.launch()