import gradio as gr
from huggingface_hub import InferenceClient
import threading
import os 

API_TOKEN = os.environ.get("API_TOKEN")

DEFAULT_NEGATIVE_INPUT = "nsfw, nude, nudity, naked, pantsless, boob, boobs, extra fingers, fewer fingers, (worst quality:1.4), (low quality:1.4), (monochrome:1.1), (bad_prompt_version2:0.8), text, jpeg, watermark, artist sign, blurry, out of frame, multiple breasts, (mutated hands and fingers:1.5 ), (long body :1.3), (mutation, poorly drawn :1.2) , black-white, bad anatomy, liquid body, liquid tongue, disfigured, malformed, mutated, anatomical nonsense, text font ui, error, malformed hands, long neck, blurred, lowers, lowres, bad anatomy, bad proportions, bad shadow, uncoordinated body, unnatural body, fused breasts, bad breasts, huge breasts, poorly drawn breasts, extra breasts, liquid breasts, heavy breasts, missing breasts, huge haunch, huge thighs, huge calf, bad hands, fused hand, missing hand, disappearing arms, disappearing thigh, disappearing calf, disappearing legs, fused ears, bad ears, poorly drawn ears, extra ears, liquid ears, heavy ears, missing ears, fused animal ears, bad animal ears, poorly drawn animal ears, extra animal ears, liquid animal ears, heavy animal ears, missing animal ears, text, ui, error, missing fingers, fused feet, bad feet, poorly drawn feet, extra feet, melting feet, fused feet, liquid feet, missing toes, bad toes, fused toes, poorly drawn toes, extra toes, fused eyes, missing eyes, liquid eyes, blurry eyes, fused lips, missing mouth, liquid mouth, expressionless, no expression, null face, liquid face, swirly face, no error correction,"

API_ENDPOINTS = {
    "Kawaii": "Ojimi/anime-kawai-diffusion",
    "Realistic": "Yntec/InsaneRealisticCVAE",
    "Pixel": "nerijs/pixel-art-xl",
    "Cartoon": "Yntec/sexyToons",
    "Sprite": "Onodofthenorth/SD_PixelArt_SpriteSheet_Generator",
    "Aesthetic": "playgroundai/playground-v2-1024px-aesthetic",
    "AestheticL": "playgroundai/playground-v2-512px-base",
    "OpenDalle": "dataautogpt3/OpenDalleV1.1",
    "Default": "gsdf/Counterfeit-V2.5",
}

CHOICES = []
CLIENTS = {}

for model_name, model_endpoint in API_ENDPOINTS.items():
    CHOICES.append(model_name)
    CLIENTS[model_name] = InferenceClient(model_endpoint, headers = { "Authorization": f"Bearer {API_TOKEN}" })

def generate(inputs, model, anti_inputs = DEFAULT_NEGATIVE_INPUT, ht = 512, wt = 512, nis = 50, gs = 7.5, seed = 42):
    
    print("MODEL CALLED: Input: " + inputs + ", Model: " + str(model))
    
    output = CLIENTS[model].text_to_image(
        inputs,
        negative_prompt = anti_inputs,
        height = int(ht),
        width = int(wt),
        num_inference_steps = int(nis),
        guidance_scale = float(gs),
        seed = int(seed)
    )

    return (output)

def cloud():
    print("[CLOUD] | Reloading all existing models...")
    for model_name, model_client in CLIENTS.items(): 
        try:
            threading.Thread(target=CLIENTS[model_name].text_to_image, args=("Preload",), kwargs={
                "negative_prompt": "",
                "height": int(512),
                "width": int(512),
                "num_inference_steps": int(1),
                "guidance_scale": float(1),
                "seed": int(0)
            }).start()
        except Exception as e:
          print(f"[PRELOAD ERROR] An error occurred: {e}")
    print("[CLOUD] | Reload successful.")

with gr.Blocks() as demo:
    with gr.Row(variant = "panel"):
        gr.Markdown("✨ A IDM space owned within Statical.")
        
    with gr.Row():
        with gr.Column():
            model = gr.Dropdown(choices = CHOICES, value = next(iter(API_ENDPOINTS)), interactive = True, label = "Model")
    with gr.Row():
        with gr.Column():
            prompt_text = gr.Textbox(label = "Your Prompt", lines = 4)
            negative_prompt_text = gr.Textbox(label = "Your Negative Prompt", lines = 4, value = DEFAULT_NEGATIVE_INPUT)
            run = gr.Button("Generate Image")
            maintain = gr.Button("☁☁️")
            height = gr.Slider( minimum = 144, maximum = 2160, value = 1024, step = 1, interactive = True, label = "Height" )
            width = gr.Slider( minimum = 144, maximum = 2160, value = 1024, step = 1, interactive = True, label = "Width" )
            nis = gr.Slider( minimum = 0, maximum = 100, value = 50, step = 1, interactive = True, label = "Steps" )
            gs = gr.Slider( minimum = 0, maximum = 100, value = 7.5, step = 0.001, interactive = True, label = "Guidance" )
            seed = gr.Slider( minimum = 0, maximum = 9007199254740991, value = 42, step = 1, interactive = True, label = "Seed" )
    with gr.Row():
        with gr.Column():
            output = gr.Image(label = "Output")
                
    run.click(generate, inputs = [prompt_text, model, negative_prompt_text, height, width, nis, gs, seed], outputs = [output], queue = False)
    maintain.click(cloud, inputs = [], outputs = [], queue = False)

demo.launch(show_api = True)