File size: 2,341 Bytes
9cc2d55
 
b855287
636bd73
b855287
a0e2cb7
b855287
a0e2cb7
b855287
adb6b80
 
b855287
a0e2cb7
b855287
a0e2cb7
9cc2d55
b855287
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a0e2cb7
9cc2d55
b855287
9cc2d55
b855287
 
 
9cc2d55
b855287
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9cc2d55
a0e2cb7
 
b855287
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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()