# @title chạy với giao diện Gradio (load mode từ local)

import gradio as gr
from diffusers import StableDiffusionPipeline
from googletrans import Translator

import torch

# Model initialization
def load_model(model_path):
    pipe = StableDiffusionPipeline.from_single_file(model_path, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
    return pipe.to("cuda")

model_path = r'D:\Truong_apacons_2023\EasySD\models\stable-diffusion\realisticVisionV51_v51VAE.safetensors'
pipe = load_model(model_path)

translator = Translator()

#mô tả tiêu cực
negative_prompt = "disfigured, ugly, verybadimagenegative_v1.3, ng_deepnegative_v1_75t, (ugly face:0.8),cross-eyed,sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), bad anatomy, DeepNegative, facing away, tilted head, {Multiple people}, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worstquality, low quality, normal quality, jpegartifacts, signature, watermark, username, blurry..."

def generate(prompt, mau_nuoc, render, height=512, width=768):
    height = int(height)
    width = int(width)

    if mau_nuoc:
        prompt = "Tranh Màu nước, " + prompt

    if render:
        prompt = "Octain render and Vray render and Corona render 3d model, fog effect, perfect perspective, perfect  edges, " + prompt

    # Xoá từ "realistic", "photographic" và "photorealistic"
    prompt = prompt.replace("realistic", "").replace("photographic", "").replace("photorealistic", "")

    # Translate
    translated_prompt = translator.translate(prompt).text

    # Hàm sinh ảnh
    images = pipe(
        prompt=translated_prompt,
        height=height,
        width=width,
        num_inference_steps=50,
        guidance_scale=7,
        negative_prompt=negative_prompt,
    ).images

    image = images[0]
    return image

Ví_dụ = [
    ["Chân dung chi tiết cao của một cô gái trẻ xinh đẹp với mái tóc trắng và đôi mắt sắc nét, 35 mm, 8k, sắc nét"],
    ["Cô gái xinh đẹp mặc váy màu vàng chơi guitar, đính kim tuyến lộng lẫy, hình ảnh sắc nét, không bị mờ"],
    ["Ngôi nhà màu vàng trong rừng, ánh sáng mặt trời tuyệt đẹp"],
    ["A cat playing with a ball of yarn"],
    ["Chú chó chạy xe đạp màu hồng"],
    ["Bình minh trên bờ biển"]
]






iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.inputs.Textbox(),
        gr.inputs.Checkbox(label="Màu nước"), 
        gr.inputs.Checkbox(label="Render"),
        gr.inputs.Number(label="Kích thước ảnh - Chiều cao", default=512),
        gr.inputs.Number(label="Kích thước ảnh - Chiều rộng", default=768)
    ],
    outputs="image",
    title="Sáng tạo hình ảnh",
    description="Mô tả về chương trình<br>Hãy viết mô tả mà bạn mong muốn vào ô nhập prompt và bấm vào nút Submit",
    examples=Ví_dụ
)

iface.launch(share=True, debug=True)