import os

# os.system('pip install paddlepaddle') # for cpu enviroment
os.system('pip install paddlepaddle-gpu==2.4.2.post117 -f https://www.paddlepaddle.org.cn/whl/linux/mkl/avx/stable.html')
os.system('pip install paddlenlp>=2.5.2')
os.system('pip install ppdiffusers>=0.14.0')

import gradio as gr
from ppdiffusers import StableDiffusionPipeline

# 通过git获取仓库模型，并读取
# import git
# 获取模型参数
# repo = git.Repo.clone_from(url='https://huggingface.co/Liyulingyue/Neolle_Face_Generator', to_path="./dream_outputs")
# 加载模型
# model_path = "dream_outputs"
# pipe = StableDiffusionPipeline.from_pretrained(model_path)

pipe = StableDiffusionPipeline.from_pretrained("Liyulingyue/Neolle_Face_Generator", from_hf_hub=True)

def generate_images(prompt, num_inference_steps, guidance_scale):
    # num_inference_steps to number
    try:
        infer_steps = int(num_inference_steps)
    except:
        infer_steps = 50

    # guidance_scale to number
    try:
        gui_scale = float(guidance_scale)
    except:
        gui_scale = 7.5
    
    image = pipe(prompt, num_inference_steps=infer_steps,guidance_scale=gui_scale).images[0]
    # image = os.getcwd()
    return image

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # 诺艾尔生成器
    基于 Linaqruf/anything-v3.0 训练，采用DreamBooth的技术并使用a photo of Neolle文本进行了训练。用于微调的图片共10张，均为原神角色诺艾尔，batch_size取1，学习率是5e-6，共训练1000步。

    Hugging face的CPU环境num_inference_steps=50时，大约需要运行1200s。在T4 small环境下，一张图大约需要30s到60s。

    如果推理结果包含色情内容，会返回一张纯黑图片~ 如果出现纯黑图片请重新运行

    欢迎大家从 https://huggingface.co/Liyulingyue/Neolle_Face_Generator 下载模型到本地运行, 20s即可出图, 该链接包含运行示例代码~

    ## 输入参数如下：
    - prompt：提示语
    - num_inference_steps: 推理轮次，越高越耗时，能够提高画作结果的精细程度，建议取值50，更高会需要消耗更多的时间，但效果会更好。
    - guidance_scale：训练图片的影响度，如果无法满足提示词描述的场景，可以降低该值，建议取值50。

    ## 推荐的提示词示例：
    - Noelle with dark hair, beautiful eyes
    - Noelle, 20 years old
    - Noelle with glasses
    - Noelle with sunglasses
    - Noelle playing basketball
    - Noelle with cat ears, blue hair
    """)
    gr.Interface(fn=generate_images, 
                 inputs=["text","text","text"], 
                 outputs="image")

if __name__ == "__main__":
    demo.launch()