import gradio as gr
import numpy as np
import random
import torch
import spaces
from diffusers import DiffusionPipeline
import requests

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v140-sdxl"
pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1536

# Translation function
@spaces.GPU
def translate_albanian_to_english(text):
    if not text.strip():
        return ""
    for attempt in range(2):
        try:
            response = requests.post(
                "https://hal1993-mdftranslation1234567890abcdef1234567890-fc073a6.hf.space/v1/translate",
                json={"from_language": "sq", "to_language": "en", "input_text": text},
                headers={"accept": "application/json", "Content-Type": "application/json"},
                timeout=5
            )
            response.raise_for_status()
            translated = response.json().get("translate", "")
            return translated
        except Exception as e:
            if attempt == 1:
                raise gr.Error(f"Përkthimi dështoi: {str(e)}")
    raise gr.Error("Përkthimi dështoi. Ju lutem provoni përsëri.")

# Aspect ratio function
def update_aspect_ratio(ratio):
    if ratio == "1:1":
        return 1024, 1024
    elif ratio == "9:16":
        return 576, 1024
    elif ratio == "16:9":
        return 1024, 576
    return 1024, 1024

@spaces.GPU(duration=120)
def infer(prompt, width, height, progress=gr.Progress(track_tqdm=True)):
    # Translate prompt
    final_prompt = translate_albanian_to_english(prompt.strip()) if prompt.strip() else ""
    final_prompt = f"score_9, score_8_up, score_7_up, source_anime, {final_prompt}"

    negative_prompt = "(low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn, (deformed | distorted | disfigured:1.3), bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers:1.4, disconnected limbs, blurry, amputation"

    seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=final_prompt,
        negative_prompt=negative_prompt,
        guidance_scale=7,
        num_inference_steps=60,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    return image

def create_demo():
    with gr.Blocks() as demo:
        # CSS for 320px gap and download button scaling
        gr.HTML("""
        <style>
        body::before {
            content: "";
            display: block;
            height: 320px;
            background-color: var(--body-background-fill);
        }
        button[aria-label="Fullscreen"], button[aria-label="Fullscreen"]:hover {
            display: none !important;
            visibility: hidden !important;
            opacity: 0 !important;
            pointer-events: none !important;
        }
        button[aria-label="Share"], button[aria-label="Share"]:hover {
            display: none !important;
        }
        button[aria-label="Download"] {
            transform: scale(3);
            transform-origin: top right;
            margin: 0 !important;
            padding: 6px !important;
        }
        </style>
        """)

        gr.Markdown("# Krijo Anime")
        gr.Markdown("Gjenero anime nga përshkrimi yt me fuqinë e inteligjencës artificiale.")

        with gr.Column():
            prompt = gr.Textbox(
                label="Përshkrimi",
                placeholder="Shkruani përshkrimin këtu",
                lines=3
            )
            aspect_ratio = gr.Radio(
                label="Raporti i fotos",
                choices=["9:16", "1:1", "16:9"],
                value="1:1"
            )
            generate_button = gr.Button(value="Gjenero")
            result_image = gr.Image(
                label="Imazhi i Gjeneruar",
                interactive=False
            )

            # Hidden sliders for width and height
            width_slider = gr.Slider(
                value=1024,
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=8,
                visible=False
            )
            height_slider = gr.Slider(
                value=1024,
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=8,
                visible=False
            )

            # Update hidden sliders based on aspect ratio
            aspect_ratio.change(
                fn=update_aspect_ratio,
                inputs=[aspect_ratio],
                outputs=[width_slider, height_slider],
                queue=False
            )

            # Bind the generate button
            generate_button.click(
                fn=infer,
                inputs=[prompt, width_slider, height_slider],
                outputs=[result_image],
                show_progress="full"
            )

    return demo

if __name__ == "__main__":
    print(f"Gradio version: {gr.__version__}")
    app = create_demo()
    app.queue(max_size=12).launch(server_name='0.0.0.0')