import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch
from langs import LANGS
import os

import numpy as np
import modin.pandas as pd
from PIL import Image
from diffusers import DiffusionPipeline #EulerDiscreteScheduler

device1 = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", safety_checker=None)
pipe = pipe.to(device1)

pipe1 = pipeline('text-generation', model='RamAnanth1/distilgpt2-sd-prompts')

TASK = "translation"
CKPT = "facebook/nllb-200-distilled-600M"

model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
tokenizer = AutoTokenizer.from_pretrained(CKPT)

device = 0 if torch.cuda.is_available() else -1


def translate(text):
    """
    Translate the text from source lang to target lang
    """
    src_lang = "zho-Hans"
    tgt_lang = "eng_Latn"
    max_length = 400
    
    translation_pipeline = pipeline(TASK,
                                    model=model,
                                    tokenizer=tokenizer,
                                    src_lang=src_lang,
                                    tgt_lang=tgt_lang,
                                    max_length=max_length,
                                    device=device)

    result = translation_pipeline(text)
    return result[0]['translation_text']

#prompt

stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
clip_interrogator_2 = gr.Blocks.load(name="spaces/fffiloni/CLIP-Interrogator-2")

def get_images(prompt):
    gallery_dir = stable_diffusion(prompt, fn_index=2)    
    img_results = [os.path.join(gallery_dir, img) for img in os.listdir(gallery_dir)]  
    return img_results[0]

def get_new_prompt(img, mode):
    interrogate = clip_interrogator_2(img, mode, 12, api_name="clipi2")
    return interrogate

def infer(input):
    prompt = pipe1(input+',', num_return_sequences=1)[0]["generated_text"]
    img = get_images(prompt) 
    result = get_new_prompt(img, 'fast')    
    return result[0]

#stable diffusion

def genie (prompt, negative_prompt, scale, steps, seed):
     generator = torch.Generator(device=device1).manual_seed(seed)
     images = pipe(prompt, negative_prompt=negative_prompt, width=768, height=768, num_inference_steps=steps, guidance_scale=scale, num_images_per_prompt=1, generator=generator).images[0]
     return images


with gr.Blocks() as demo:
    gr.Markdown(
            """ # <center>🥳💬💕 - Stable Diffusion All in One，随时随地，想画就画！</center>
            
            ### <center>🤖 - 让有人文关怀的AI造福每一个人！AI向善，文明璀璨！山水AI - Enable the future！</center>
            
        """
    )
    
    with gr.Row().style(equal_height=True):
        inp1 = gr.Textbox(label = "请输入中文提示词", lines=1)
        button = gr.Button("翻译成英文吧")
        out1 = gr.Textbox(label = "翻译后的英文提示词", lines=1)
    
    button.click(translate, [inp1], [out1])
    
    with gr.Row().style(equal_height=True):
        inp2 = out1
        btn1 = gr.Button("让您的提示词更详细一些吧")
        out2 = gr.Textbox(label = "更详细的英文提示词(可编辑此内容)", lines=1)
    
    btn1.click(infer, [inp2], [out2])

    with gr.Row().style(equal_height=True):
        inp3 = out2
        inp4 = gr.Textbox(label="反向关键词：删除不想要的图片特点(选填)", placeholder = "low quality, illustration", lines=1)
        inp5 = gr.Slider(1, 25, 10, label="提示词的符合度", info="数值越大，图片越符合提示词，但同时画质会下降")
        inp6 = gr.Slider(30, maximum=100, value=50, step=1, label="生成图片的步数", info="数值越大，图片越清晰，但运行时间会变长")
        inp7 = gr.Slider(minimum=1, step=1, maximum=999999999999999999, randomize=True, label="seed", info="生成一个随机数")

    with gr.Row().style(equal_height=True):
        btn2 = gr.Button("开始生成图片吧")
        out3 = gr.Image(label="为您生成的专属图片")

    btn2.click(genie, [inp3, inp4, inp5, inp6, inp7], [out3])
    
    gr.Markdown(
            """ ### <center>注意❗：请不要输入或生成会对个人以及组织造成侵害的内容，此程序仅供科研、学习及娱乐使用。用户输入或生成的内容与程序开发者无关，请自觉合法合规使用，违反者一切后果自负。</center>
            
            ### <center>Model by [Stable Diffusion](https://github.com/Stability-AI/stablediffusion). Thanks to [Manjushri](https://huggingface.co/Manjushri).</center>
            
        """
    )
        
    gr.HTML('''
        <div class="footer">
                    <p>🎶🖼️🎡 - It’s the intersection of technology and liberal arts that makes our hearts sing. - Steve Jobs
                    </p>
        </div>
    ''')     

demo.queue().launch(show_error=True)