import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
from nextstep.models.pipeline_nextstep import NextStepPipeline
from spaces import GPU

HF_HUB = "stepfun-ai/NextStep-1-Large"

# Load model & tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_HUB, trust_remote_code=True)
model = AutoModel.from_pretrained(HF_HUB, trust_remote_code=True)
pipeline = NextStepPipeline(tokenizer=tokenizer, model=model).to(device="cuda", dtype=torch.bfloat16)

IMG_SIZE = 512

@spaces.GPU
def generate(main_prompt, positive_prompt, negative_prompt):
    image = pipeline.generate_image(
        main_prompt,
        hw=(IMG_SIZE, IMG_SIZE),
        num_images_per_caption=1,
        positive_prompt=positive_prompt,
        negative_prompt=negative_prompt,
        cfg=7.5,
        cfg_img=1.0,
        cfg_schedule="constant",
        use_norm=False,
        num_sampling_steps=28,
        timesteps_shift=1.0,
        seed=3407,
    )[0]
    return image

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Main Prompt", lines=2, placeholder="Describe your dream image..."),
        gr.Textbox(label="Positive Prompt", lines=1, placeholder="Extra quality boosters..."),
        gr.Textbox(label="Negative Prompt", lines=1, placeholder="Things to avoid..."),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="NextStep Image Generator",
    description="Simple Gradio demo for NextStep-1-Large with editable prompts.",
    examples=[
        [
            'A realistic photograph of a wall with "NextStep-1.1 is coming" prominently displayed',
            'masterpiece, film grained, best quality.',
            'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.'
        ],
        [
            'A cinematic shot of a futuristic city at sunset',
            'ultra-detailed, 8k, volumetric lighting, masterpiece',
            'blurry, lowres, noise, distortion'
        ],
        [
            'A fantasy painting of a dragon flying over a castle',
            'high detail, vibrant colors, epic composition',
            'washed out colors, bad anatomy, low detail'
        ],
    ]
)

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