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("Completed")

@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 Image Generator")
    gr.Markdown("Generate high-quality anime-style images using the NewBie model.")
    
    with gr.Row(variant="panel"):
        with gr.Column(scale=2):
            prompt_newbie = gr.Textbox(
                label="Prompt",
                value=newbie_prompt,
                lines=10,
                placeholder="Enter the prompt here..."
            )
            negative_prompt_newbie = gr.Textbox(
                label="Negative Prompt",
                value="low quality, bad quality, blurry, low resolution, deformed, ugly, bad anatomy",
                lines=3,
                placeholder="Enter the elements that should not appear here..."
            )
            
            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 is random", 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()