import gradio as gr
from ppdiffusers import StableDiffusionPipeline
import os
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)

def generate_images(prompt="Neolle", num_inference_steps=50, guidance_scale=7.5):
    image = pipe(prompt, num_inference_steps=int(num_inference_steps),guidance_scale=float(guidance_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。

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

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

    ## 输入参数如下：
    - prompt：提示语
    - num_inference_steps: 推理轮次，越高越耗时，能够提高画作结果的精细程度，默认50
    - guidance_scale：训练图片的影响度，如果无法满足提示词描述的场景，可以降低该值，默认7.5

    ## 推荐的提示词示例：
    - Noelle with glasses
    - Noelle with sunglasses
    - Noelle with dark hair, beautiful eyes
    - Noelle, 20 years old
    - 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()