import gradio as gr
import requests
import os
from PIL import Image
from io import BytesIO
from tqdm import tqdm
import time
import numpy as np
import potrace

# Defining the repository information and the trigger word
repo = "artificialguybr/LineAniRedmond-LinearMangaSDXL-V2"
trigger_word = "lineart,LineAniAF,"

# Hard set prompt template
hard_set_prompt = (
    "minimalist black and white single-line art illustration of [subject], created in one continuous, unbroken stroke from start to finish. "
    "The design uses true black and true white values only, with no gradients or gray tones. The artwork emphasizes clean, connected, and flowing lines, "
    "resulting in a sleek, modern aesthetic. The [subject] is depicted in a dynamic and recognizable pose, optimized for artistic purposes and ready for 3D printing applications."
)

constraints = (
    "MUST: be only black and white, must be single-line art, must never use gray values/shading/tint or hue of any kind, only true black and true white. "
    "NEVER use anything besides a single black line of varying widths as necessary to achieve a result of the user's request."
)

def generate_image(prompt):
    print("Generating image with prompt:", prompt)
    api_url = f"https://api-inference.huggingface.co/models/{repo}"
    token = os.getenv("API_TOKEN")
    headers = {
        "Authorization": f"Bearer {token}"
    }
    
    # Incorporate the hard set prompt and constraints
    full_prompt = (
        f"{hard_set_prompt.replace('[subject]', prompt)} {constraints} {trigger_word}"
    )
    payload = {
        "inputs": full_prompt,
        "parameters": {
            "negative_prompt": "(non-minimalist, best quality, high quality, normal quality, hires, details, anything but basic 1-line-art inferred from user submitted text, overexposed, underexposed, grayscale, bw, bad photo, bad photography, bad art:1.4), (watermark, signature, text font, username, error, logo, words, letters, digits, autograph, trademark, name:1.2), (blur, blurry, grainy), morbid, ugly, asymmetrical, mutated malformed, mutilated, poorly lit, bad shadow, draft, cropped, out of frame, cut off, censored, jpeg artifacts, out of focus, glitch, duplicate, (airbrushed, cartoon, anime, semi-realistic, cgi, render, blender, digital art, manga, amateur:1.3), , :1.3)",
            "num_inference_steps": 30,
            "scheduler": "DPMSolverMultistepScheduler"
        },
    }

    error_count = 0
    pbar = tqdm(total=None, desc="Loading model")
    while True:
        print("Sending request to API...")
        response = requests.post(api_url, headers=headers, json=payload)
        print("API response status code:", response.status_code)
        if response.status_code == 200:
            print("Image generation successful!")
            # Open the image and convert it to grayscale
            # Convert the grayscale image to a binary image
            image = image.point(lambda p: 255 if p > 128 else 0, mode='1')
            # Convert the PIL image to a numpy array
            array = np.array(image)
            # Invert the array for potrace (black on white)
            array = 255 - array
            # Trace the bitmap to SVG
            bitmap = potrace.Bitmap(array)
            path = bitmap.trace()
            # Generate SVG data
            svg_data = path.to_svg()
            return svg_data
    
        elif response.status_code == 503:
            time.sleep(1)
            pbar.update(1)
        elif response.status_code == 500 and error_count < 5:
            time.sleep(1)
            error_count += 1
        else:
            print("API Error:", response.status_code)
            raise Exception(f"API Error: {response.status_code}")

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Describe the subject here..."),
    outputs="text",  # SVG data as text
    title="LineArt XL Image Generator",
    description=(
        "Powered by the generous GPU time from Redmond.AI, this LORA, fine-tuned on SD XL 1.0, excels at creating Lineart-themed images across a wide range of subjects. "
        "Optimized for 1024x1024 resolution, it incorporates the specific tag 'lineart,LineAniAF' directly in the HF Space for ease of use. "
        "If you appreciate this model and wish to support, consider a donation via [Patreon](https://www.patreon.com/user?u=81570187) or [Ko-fi](https://ko-fi.com/artificialguybr). "
        "Stay updated on new models by following on [Twitter](https://twitter.com/artificialguybr)."
    ),
    examples=[
        ["a soaring eagle"], 
        ["a running cheetah"], 
        ["a majestic oak tree"]
    ]
)

print("Launching Gradio interface...")
iface.launch()