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

def generate_image(prompt, negative_prompt, scheduler, steps, width, height, cfg, restore_faces, seed):
    request_time = datetime.datetime.now()
    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))

           # Log the information together
           print(f"Request time: {request_time}\n"
                  f"Prompt: {prompt}\n"
                  f"Negative Prompt: {negative_prompt}\n"
                  f"Seed: {seed}\n"
                  f"Res(width x height): {width} x {height}\n"
                  f"Image URL: {image_url}")

           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=1024, label="Width"),
        gr.components.Slider(minimum=512, maximum=1600, value=1024, 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](https://civitai.com/models/87288/freedomredmond) with example prompts in CIVITAI and [here in HF](https://huggingface.co/artificialguybr/freedom).

## 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.

## Due to high demand, the generations are taking longer. Please wait and your image will be ready.
""",
    allow_flagging='never'
)

#Adding queue
iface.queue(concurrency_count=12)

# Launch the app
iface.launch()