import gradio as gr
import os
import re
#from tempfile import NamedTemporaryFile
import numpy as np
#import spaces
import random
import string
from diffusers import StableDiffusion3Pipeline
import torch
from pathos.multiprocessing import ProcessingPool as ProcessPoolExecutor
import requests
from lxml.html import fromstring

pool = ProcessPoolExecutor(100)
pool.__enter__()

#model_id = "runwayml/stable-diffusion-v1-5"
model_id = "stabilityai/stable-diffusion-3-medium-diffusers"

device = "cuda" if torch.cuda.is_available() else "cpu"
if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True, token=os.getenv('hf_token'))
    pipe = pipe.to(device)
else: 
    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, use_safetensors=True, token=os.getenv('hf_token'))
    pipe = pipe.to(device)

def translate(text,lang):
    user_agents = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 13_1) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15'
    ]
    url = f'https://www.google.com/search?q=translate to {lang}: {text}'
    print(url)
    resp = requests.get(
        url = url,
        headers = {
            'User-Agent': random.choice(user_agents)
        }
    )
    print(resp)
    content = resp.content
    html = fromstring(content)
    rslt = html.xpath('//pre[@aria-label="Translated text"]/span')
    translated = rslt[0].text.strip()
    ret = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', translated)).lower().strip()
    print(ret)
    return ret

def generate_random_string(length):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))

#@spaces.GPU
def Piper(english_prompt):
    return pipe(
        english_prompt,
        height=256,
        width=256,
        negative_prompt="",
        num_inference_steps=200,
        guidance_scale=10
    )

def infer(prompt):
    name = generate_random_string(12)+".png"
    english_prompt = f'REAL {translate(prompt,"english").upper()}:'
    print(f'Final prompt: {english_prompt}')
    image = Piper(english_prompt).images[0].save(name)
    return name

css="""
#col-container {
    margin: 0 auto;
    max-width: 14cm;
}
#image-container {
    aspect-ratio: 1 / 1;
}
"""

if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
            # Image Generator
            Currently running on {power_device}.
        """)
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        result = gr.Image(elem_id="image-container", label="Result", show_label=False, type='filepath')
    run_button.click(
        fn = infer,
        inputs = [prompt],
        outputs = [result]
    )

demo.queue().launch()