import gradio as gr
import random
import os


default_negative_prompt = (
    "Extra limbs, Extra fingers or toes, Disfigured face, Distorted hands, Mutated body parts, "
    "Missing limbs, Asymmetrical features, Blurry face, Poor anatomy, Incorrect proportions, Crooked eyes, "
    "Deformed hands or fingers, Double face, Unrealistic skin texture, Overly smooth skin, Poor lighting, "
    "Cartoonish appearance, Plastic look, Grainy, Unnatural expressions, Crossed eyes, Mutated clothing, "
    "Artifacts, Uncanny valley, Doll-like features, Bad symmetry, Uneven skin tones, Extra teeth, "
    "Unrealistic hair texture, Dark shadows on face, Overexposed face, Cluttered background, Text, watermark, "
    "or signature, Over-processed, Low quality, Blurry, Low resolution, Pixelated, Oversaturated, Too dark, Overexposed, Poor lighting, "
    "Unclear, Text or watermark, Distorted faces, Extra limbs or fingers, Disfigured, Grainy, Overly stylized, Cartoonish, "
    "Unrealistic anatomy, Bad proportions, Unrealistic colors, Artificial look, Background noise, Unwanted objects, Repetitive patterns, Artifacting, Abstract shapes, Out of focus"
)

model = gr.load("models/Purz/face-projection")

def generate_image(text, seed, width, height, guidance_scale, num_inference_steps, negative_prompt):
    if seed is not None:
        random.seed(seed)

    if text in [example[0] for example in examples]:
        print(f"Using example: {text}")

    result_image = model(text, negative_prompt=negative_prompt)
    
    print(f"Width: {width}, Height: {height}, Guidance Scale: {guidance_scale}, Inference Steps: {num_inference_steps}, Negative Prompt: {negative_prompt}")
    
    return result_image

def randomize_parameters():
    seed = random.randint(0, 999999)
    width = random.randint(512, 2048)
    height = random.randint(512, 2048)
    guidance_scale = round(random.uniform(0.1, 20.0), 1)
    num_inference_steps = random.randint(1, 40)
    
    return seed, width, height, guidance_scale, num_inference_steps

examples = [
    ["Create an image of a fierce humanoid cat warrior standing in full view. The character has a muscular, agile build with feline features, including sharp ears, whiskers, and a tail. Their fur is a sleek combination of dark gray and silver, with patterns resembling a wild cat. The warrior wears lightweight, yet protective armor made of metal and leather, designed for agility in battle. The armor has intricate detailing, with claws and paw-like motifs, and it covers key areas such as the shoulders, chest, and legs. The warrior holds a large, double-edged sword in one hand, and a small shield with a paw print insignia in the other. Their fierce eyes glow with determination, and their stance is confident, ready for combat. The background is a simple, blurred forest or battlefield, ensuring the focus remains on the warrior.", *randomize_parameters()],
    ["Create an epic scene of a Warhammer Sisterhood army preparing for battle. The army consists of fierce, armored female warriors dressed in intricate power armor, adorned with symbols of their faith and dedication to their cause. They carry a mix of weapons, including bolters, power swords, and flamers, all gleaming with metallic sheen. The warriors' armor is a striking combination of white, gold, and red, with ornate detailing and religious icons emblazoned across their chests and shoulders. The backdrop shows a war-torn battlefield with the ruins of a grand temple or city in the distance, smoke rising into the sky. Their leader, standing at the forefront, holds a mighty relic weapon and commands the troops with unwavering confidence. The scene exudes strength, faith, and determination as the Sisterhood prepares to face their enemies in a brutal war.", *randomize_parameters()],
    ["Create an intense scene of a future robot war, set in a dystopian cityscape. The foreground shows heavily armored robots, equipped with advanced weaponry and glowing energy cores. Some robots have sleek, metallic designs with sharp, angular features, while others are more rugged with battle-worn exteriors. The robots are engaged in fierce combat, firing laser beams and launching missiles at each other. Explosions and sparks fly in the background, with a smoky, debris-filled atmosphere. The city is in ruins, with tall, shattered skyscrapers and wreckage scattered across the streets. The sky above is dark, with futuristic aircraft zooming by in the distance.", *randomize_parameters()],
    ["Create an image of a beautiful female model wearing a sporty outfit consisting of a stylish sports bra and leggings. She is standing confidently, showcasing a fit and athletic physique. The sports bra is sleek and modern, paired with matching leggings that accentuate her shape. The leggings are designed with a trendy pattern or color, and she is wearing bright pink sneakers that stand out. Her hair is styled in a natural, athletic look, and the background is simple and minimalistic to highlight her pose and outfit.",*randomize_parameters()],
]

interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Type here your imagination:", placeholder="Type or click an example..."),
        gr.Slider(label="Seed", minimum=0, maximum=999999, step=1),
        gr.Slider(label="Width", minimum=512, maximum=2048, step=64, value=1024),
        gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1024),
        gr.Slider(label="Guidance Scale", minimum=0.1, maximum=20.0, step=0.1, value=3.0),
        gr.Slider(label="Number of inference steps", minimum=1, maximum=40, step=1, value=28),
        gr.Dropdown(
            label="Negative Prompt",
            choices=[
                "Default Negative Prompt",
                "None",
            ],
            value="Default Negative Prompt",
            description="Choose a negative prompt to apply to the image generation."
        ),
    ],
    outputs=gr.Image(label="Generated Image"),
    examples=examples,
    theme="NoCrypt/miku",
    description="Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
)

def get_negative_prompt(negative_prompt_choice):
    if negative_prompt_choice == "Default Negative Prompt":
        return default_negative_prompt
    else:
        return ""


def generate_image_with_negative_prompt(text, seed, width, height, guidance_scale, num_inference_steps, negative_prompt_choice):
    negative_prompt = get_negative_prompt(negative_prompt_choice)
    return generate_image(text, seed, width, height, guidance_scale, num_inference_steps, negative_prompt)

interface.fn = generate_image_with_negative_prompt

interface.launch()