import gradio as gr
import spaces
import torch
from diffusers import NewbiePipeline
from transformers import AutoModel
import random
import warnings

warnings.filterwarnings("ignore")

model_path = "Disty0/NewBie-image-Exp0.1-Diffusers"  # NewBie-AI/NewBie-image-Exp0.1

print("正在加载 Text Encoder...")
text_encoder_2 = AutoModel.from_pretrained(
    model_path, 
    subfolder="text_encoder_2", 
    trust_remote_code=True, 
    torch_dtype=torch.bfloat16
)

print("正在加载 Pipeline...")
pipe_newbie = NewbiePipeline.from_pretrained(
    model_path, 
    text_encoder_2=text_encoder_2, 
    torch_dtype=torch.bfloat16
)
pipe_newbie.to("cuda")

del text_encoder_2
torch.cuda.empty_cache()
print("模型加载完成。")

@spaces.GPU()
def generate_image_newbie(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    if seed < 0:
        seed = random.randint(0, 2**32 - 1)

    generator = torch.Generator("cuda").manual_seed(int(seed))

    image = pipe_newbie(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]

    return image, seed

# 默认提示词
newbie_prompt = """<character_1>
<n>$character_1$</n>
<gender>1girl, solo</gender>
<appearance>blonde_hair, long_hair</appearance>
<clothing>large_hat, white_hat, white_blouse, puffy_sleeves, shoulder_cutout, black_skirt, shirt_tucked_in, socks, shoes</clothing>
<expression>looking_up</expression>
<action>sitting, reclining, arm_support, from_side, cowboy_shot, wide_shot</action>
<position>center</position>
</character_1>

<general_tags>
<count>1girl</count>
<artists> (kazutake hazano:1.2), (kazutake hazano:0.5), (onineko:0.8), (r17329 illu:0.2), (ma1ma1helmes b illu:0.2)</artists>
<style>masterpiece, best_quality, high_resolution, detailed</style>
<background>detailed_background, scenery, detailed_background</background>
<atmosphere>cheerful</atmosphere>
<lighting>dynamic_angle, depth_of_field, high_contrast, colorful, detailed_light, light_leaks, beautiful_detailed_glow, best_shadow, shiny_skin, cinematic_lighting, ray_tracing, from_above, female_focus, close-up, dutch_angle, blue_archive</lighting>
<quality>very_aesthetic, masterpiece, no_text</quality>
<objects>bag</objects>
<other>2024_year</other>
</general_tags>"""

with gr.Blocks(title="NewBie 图像生成器") as demo:
    gr.Markdown("# NewBie 图像生成器 (NewBie Image Generator)")
    gr.Markdown("使用 NewBie 模型生成高质量动漫风格图像。")
    
    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            prompt_newbie = gr.Textbox(
                label="提示词 (Prompt)",
                value=newbie_prompt,
                lines=10,
                placeholder="在此输入生成提示词..."
            )
            negative_prompt_newbie = gr.Textbox(
                label="负面提示词 (Negative Prompt)",
                value="low quality, bad quality, blurry, low resolution, deformed, ugly, bad anatomy",
                lines=3,
                placeholder="在此输入不需要出现的元素..."
            )
            
            with gr.Row():
                height_newbie = gr.Slider(label="图片高度 (Height)", minimum=512, maximum=2048, step=64, value=1024)
                width_newbie = gr.Slider(label="图片宽度 (Width)", minimum=512, maximum=2048, step=64, value=1024)
            
            with gr.Row():
                steps_newbie = gr.Slider(label="推理步数 (Inference Steps)", minimum=1, maximum=100, step=1, value=28)
                guidance_scale_newbie = gr.Slider(label="引导系数 (Guidance Scale)", minimum=1.0, maximum=20.0, step=0.1, value=3.5)
            
            seed_newbie = gr.Number(label="随机种子 (Seed, -1 为随机)", value=-1, precision=0)
            
            generate_btn_newbie = gr.Button("生成图片 (Generate)", variant="primary")

        with gr.Column(scale=1):
            image_output_newbie = gr.Image(label="生成结果 (Output)", format="png", interactive=False)
            used_seed_newbie = gr.Number(label="使用的种子 (Used Seed)", interactive=False)

    generate_btn_newbie.click(
        fn=generate_image_newbie,
        inputs=[
            prompt_newbie, 
            negative_prompt_newbie, 
            height_newbie, 
            width_newbie, 
            steps_newbie, 
            guidance_scale_newbie, 
            seed_newbie
        ],
        outputs=[image_output_newbie, used_seed_newbie]
    )

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