import gradio as gr
import requests
import json
import PIL.Image
from io import BytesIO
import os
import random

def generate_image(prompt, negative_prompt, scheduler, steps, width, height, restore_faces, seed, cfg):
    restore_faces = bool(restore_faces)
    print(f"restore_faces: {restore_faces}, type: {type(restore_faces)}")
    # Define the API endpoint
    apiUrl = os.getenv("API_URL")
    # Define the request headers
    headers = {
        "Content-Type": "application/json",
        "token": os.getenv("API_TOKEN")
    }

    # Define the request body
    body = {
        "mode": "url",
        "model": "Freedom.safetensors",
        "tiling": False,
        "batch_size": 1,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "seed":random.randint(0, 999999999),
        "scheduler": scheduler,
        "n_iter": 1,
        "steps": steps,
        "cfg": cfg,
        "offset_noise": 0.0,
        "width": width,
        "height": height,
        "clip_skip": 1,
        "vae": "vae-ft-mse-840000-ema-pruned.ckpt",
        "restore_faces": restore_faces,
        "fr_model": "CodeFormer",
        "codeformer_weight": 0.5,
        "enable_hr": False,
        "denoising_strength": 0.75,
        "hr_scale": 2,
        "hr_upscale": "None",
        "img2img_ref_img_type": "piece",
        "img2img_resize_mode": 0,
        "img2img_denoising_strength": 0.75,
    }

    # Send the request
    response = requests.post(apiUrl, headers=headers, data=json.dumps(body), verify=False)
    # Print the response body if the status code is not 200
    if response.status_code != 200:
        print(response.text)

    # Check the response status
    if response.status_code == 200:
        
       # Get the image URL from the response
       response_json = response.json()
       if 'results' in response_json and isinstance(response_json['results'], list) and len(response_json['results']) > 0:
           image_url = response_json['results'][0]

           # Get the image from the URL           
           image_response = requests.get(image_url)
           image = PIL.Image.open(BytesIO(image_response.content))

           return image
       else:
           raise Exception("Unexpected API response format")
    else:
        raise Exception("API request failed with status code " + str(response.status_code))

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_image, 
    inputs=[
        gr.components.Textbox(label="Prompt"),
        gr.components.Textbox(value="ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft", label="Negative Prompt"),
        gr.components.Dropdown(choices=[
            "Euler a",
            "Euler",
            "LMS",
            "Heun",
            "DPM2",
            "DPM2 a",
            "DPM++ 2S a",
            "DPM++ 2M",
            "DPM++ SDE",
            "DPM fast",
            "DPM adaptive",
            "LMS Karras",
            "DPM2 Karras",
            "DPM2 a Karras",
            "DPM++ 2S a Karras",
            "DPM++ 2M Karras",
            "DPM++ SDE Karras",
            "DDIM",
            "PLMS"
        ], label="Scheduler", value="DPM++ SDE Karras"),
        gr.components.Slider(minimum=10, maximum=100, step=1.0,value=30, label="Steps"),
        gr.components.Slider(minimum=512, maximum=1600, value=768, label="Width"),
        gr.components.Slider(minimum=512, maximum=1600, value=768, label="Height"),
        gr.components.Slider(minimum=4, maximum=12, step=0.5, value=7.0, label="CFG"),
        gr.inputs.Checkbox(label="Restore Faces", default=False),
    ], 
    outputs=gr.components.Image(),
    title="Freedom.Redmond Demonstration",
    description = """
##Finetuned model of SD 2.1 768X produced by [@artificialguybr](https://twitter.com/artificialguybr).

## Resources
- The weights were released [here](LINK_TO_WEIGHTS).
- You can find example prompts [here](LINK_TO_EXAMPLE_PROMPTS).

## Demonstration
This demonstration is running on the [makeai.run API](https://www.makeai.run/).

## Acknowledgements
Thanks to [Redmond.ai](https://redmond.ai/) for providing GPU Time and sponsoring this model.
""",
    allow_flagging='never'
)

# Launch the app
iface.launch()