import gradio as gr
import requests
import json
import PIL.Image
from io import BytesIO
import os
def generate_image(prompt):
    # Define the API endpoint
    apiUrl = os.getenv("API_URL")
    token = os.getenv("API_TOKEN")
    # Define the request headers
    headers = {
        "Content-Type": "application/json",
        "token": "token"
    }

    # Define the request body
    body = {
        "mode": "url",
        "model": "AOM3A1B_orangemixs.safetensors",
        "tiling": False,
        "batch_size": 1,
        "prompt": prompt,
        "negative_prompt": "worst quality, lowres",
        "seed": 1234,
        "scheduler": "Euler a",
        "n_iter": 1,
        "steps": 25,
        "cfg": 11.0,
        "offset_noise": 0.0,
        "width": 512,
        "height": 512,
        "clip_skip": 1,
        "loras": [{"name": "", "strength": 1.0}],
        "embeddings": [{"name": "", "strength": 1.0}],
        "vae": "vae-ft-mse-840000-ema-pruned.ckpt",
        "restore_faces": False,
        "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,
        "controlnet_enabled": False,
        "controlnet_ref_img_type": "piece",
        "controlnet_guessmode": False,
        "controlnet_module": "canny",
        "controlnet_model": "control_v11p_sd15_softedge",
        "controlnet_weight": 1,
        "controlnet_guidance_start": 0,
        "controlnet_guidance_end": 1,
        "controlnet_ref_img_url": "https://upload.wikimedia.org/wikipedia/commons/d/d1/Image_not_available.png",
        "controlnet_mask": [],
        "controlnet_resize_mode": "Scale to Fit (Inner Fit)",
        "controlnet_lowvram": False,
        "controlnet_processor_res": 512,
        "controlnet_threshold_a": 100,
        "controlnet_threshold_b": 200
    }

    # Send the request
    response = requests.post(apiUrl, headers=headers, data=json.dumps(body))

    # Check the response status
    if response.status_code == 200:
        # Get the image URL from the response
        image_url = json.loads(response.text)['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("API request failed with status code " + str(response.status_code))

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_image, 
    inputs="text", 
    outputs="image",
    title="Freedom Demonstration",
    description="Enter a prompt and generate an image.",
    allow_flagging=False
)

# Launch the app
iface.launch()