import gradio as gr import numpy as np import spaces from PIL import Image import torch from torch.amp import autocast from transformers import AutoTokenizer, AutoModel from models.gen_pipeline import NextStepPipeline HF_HUB = "stepfun-ai/NextStep-1-Large" device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = AutoTokenizer.from_pretrained(HF_HUB, local_files_only=False, trust_remote_code=True) model = AutoModel.from_pretrained( HF_HUB, local_files_only=False, trust_remote_code=True, torch_dtype=torch.bfloat16, ).to(device) pipeline = NextStepPipeline(tokenizer=tokenizer, model=model).to(device=device, dtype=torch.bfloat16) MAX_SEED = np.iinfo(np.int16).max DEFAULT_POSITIVE_PROMPT = None DEFAULT_NEGATIVE_PROMPT = None DEFAULT_CFG = 7.5 def _ensure_pil(x): """Ensure returned image is a PIL.Image.Image.""" if isinstance(x, Image.Image): return x import numpy as np if hasattr(x, "detach"): x = x.detach().float().clamp(0, 1).cpu().numpy() if isinstance(x, np.ndarray): if x.dtype != np.uint8: x = (x * 255.0).clip(0, 255).astype(np.uint8) if x.ndim == 3 and x.shape[0] in (1, 3, 4): # CHW -> HWC x = np.moveaxis(x, 0, -1) return Image.fromarray(x) raise TypeError("Unsupported image type returned by pipeline.") def infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress): """Core inference logic without GPU decorators.""" if prompt in [None, ""]: gr.Warning("⚠️ Please enter a prompt!") return None with autocast(device_type=("cuda" if device == "cuda" else "cpu"), dtype=torch.bfloat16): imgs = pipeline.generate_image( prompt, hw=(int(height), int(width)), num_images_per_caption=1, positive_prompt=positive_prompt, negative_prompt=negative_prompt, cfg=float(cfg), cfg_img=1.0, cfg_schedule="constant", use_norm=False, num_sampling_steps=int(num_inference_steps), timesteps_shift=1.0, seed=int(seed), progress=True, ) return _ensure_pil(imgs[0]) # Tier 1: Very small images with few steps @spaces.GPU(duration=90) def infer_tiny(prompt=None, seed=0, width=512, height=512, num_inference_steps=24, cfg=DEFAULT_CFG, positive_prompt=DEFAULT_POSITIVE_PROMPT, negative_prompt=DEFAULT_NEGATIVE_PROMPT, progress=gr.Progress(track_tqdm=True)): return infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress) # Tier 2: Small to medium images with standard steps @spaces.GPU(duration=150) def infer_fast(prompt=None, seed=0, width=512, height=512, num_inference_steps=24, cfg=DEFAULT_CFG, positive_prompt=DEFAULT_POSITIVE_PROMPT, negative_prompt=DEFAULT_NEGATIVE_PROMPT, progress=gr.Progress(track_tqdm=True)): return infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress) # Tier 3: Standard generation for most common cases @spaces.GPU(duration=200) def infer_std(prompt=None, seed=0, width=512, height=512, num_inference_steps=28, cfg=DEFAULT_CFG, positive_prompt=DEFAULT_POSITIVE_PROMPT, negative_prompt=DEFAULT_NEGATIVE_PROMPT, progress=gr.Progress(track_tqdm=True)): return infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress) # Tier 4: Larger images or more steps @spaces.GPU(duration=300) def infer_long(prompt=None, seed=0, width=512, height=512, num_inference_steps=36, cfg=DEFAULT_CFG, positive_prompt=DEFAULT_POSITIVE_PROMPT, negative_prompt=DEFAULT_NEGATIVE_PROMPT, progress=gr.Progress(track_tqdm=True)): return infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress) # Tier 5: Maximum quality with many steps @spaces.GPU(duration=400) def infer_max(prompt=None, seed=0, width=512, height=512, num_inference_steps=45, cfg=DEFAULT_CFG, positive_prompt=DEFAULT_POSITIVE_PROMPT, negative_prompt=DEFAULT_NEGATIVE_PROMPT, progress=gr.Progress(track_tqdm=True)): return infer_core(prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt, progress) # Improved JS dispatcher with better calculation logic js_dispatch = """ function(width, height, steps){ const w = Number(width); const h = Number(height); const s = Number(steps); // Calculate total pixels and complexity score const pixels = w * h; const megapixels = pixels / 1000000; // Complexity score combines image size and steps // Base: ~0.5 seconds per megapixel per step const complexity = megapixels * s; let target = 'btn-std'; // Default // Select appropriate tier based on complexity if (pixels <= 256*256 && s <= 20) { // Very small images with few steps target = 'btn-tiny'; } else if (complexity < 5) { // Small images or few steps (e.g., 384x384 @ 24 steps = 3.5) target = 'btn-fast'; } else if (complexity < 8) { // Standard generation (e.g., 512x512 @ 28 steps = 7.3) target = 'btn-std'; } else if (complexity < 12) { // Larger or more steps (e.g., 512x512 @ 40 steps = 10.5) target = 'btn-long'; } else { // Maximum complexity target = 'btn-max'; } // Special cases: override based on extreme values if (s >= 45) { target = 'btn-max'; // Many steps always need more time } else if (pixels >= 512*512 && s >= 35) { target = 'btn-long'; // Large images with many steps } console.log(`Resolution: ${w}x${h}, Steps: ${s}, Complexity: ${complexity.toFixed(2)}, Selected: ${target}`); const b = document.getElementById(target); if (b) b.click(); } """ css = """ #col-container { margin: 0 auto; max-width: 800px; } /* Hide the dispatcher buttons */ #btn-tiny, #btn-fast, #btn-std, #btn-long, #btn-max { display: none !important; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("# NextStep-1-Large — Image generation") with gr.Row(): prompt = gr.Text(label="Prompt", show_label=False, max_lines=2, placeholder="Enter your prompt", container=False) run_button = gr.Button("Run", scale=0, variant="primary") cancel_button = gr.Button("Cancel", scale=0, variant="secondary") with gr.Row(): with gr.Accordion("Advanced Settings", open=True): positive_prompt = gr.Text(label="Positive Prompt", show_label=True, placeholder="Optional: add positives") negative_prompt = gr.Text(label="Negative Prompt", show_label=True, placeholder="Optional: add negatives") with gr.Row(): seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=3407) num_inference_steps = gr.Slider(label="Sampling steps", minimum=10, maximum=50, step=1, value=28) with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=512, step=64, value=512) height = gr.Slider(label="Height", minimum=256, maximum=512, step=64, value=512) cfg = gr.Slider(label="CFG (guidance scale)", minimum=0.0, maximum=20.0, step=0.5, value=DEFAULT_CFG, info="Higher = closer to text, lower = more creative") with gr.Row(): result_1 = gr.Image(label="Result", format="png", interactive=False) # Hidden dispatcher buttons with gr.Row(visible=False): btn_tiny = gr.Button(visible=False, elem_id="btn-tiny") btn_fast = gr.Button(visible=False, elem_id="btn-fast") btn_std = gr.Button(visible=False, elem_id="btn-std") btn_long = gr.Button(visible=False, elem_id="btn-long") btn_max = gr.Button(visible=False, elem_id="btn-max") examples = [ [ "Studio portrait of an elderly sailor with a weathered face, dramatic Rembrandt lighting, shallow depth of field", 101, 512, 512, 32, 7.5, "photorealistic, sharp eyes, detailed skin texture, soft rim light, 85mm lens", "over-smoothed skin, plastic look, extra limbs, watermark"], ["Isometric cozy coffee shop interior with hanging plants and warm Edison bulbs", 202, 512, 384, 30, 8.5, "isometric view, clean lines, stylized, warm ambience, detailed furniture", "text, logo, watermark, perspective distortion"], ["Ultra-wide desert canyon at golden hour with long shadows and dust in the air", 303, 512, 320, 28, 7.0, "cinematic, volumetric light, natural colors, high dynamic range", "over-saturated, haze artifacts, blown highlights"], ["Oil painting of a stormy sea with a lighthouse, thick impasto brushwork", 707, 384, 512, 34, 7.0, "textured canvas, visible brush strokes, dramatic sky, moody lighting", "smooth digital look, airbrush, neon colors"], ] gr.Examples( examples=examples, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], label="Click & Fill Examples (Exact Size)", ) # Wire up the dispatcher buttons to their respective functions ev_tiny = btn_tiny.click(infer_tiny, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], outputs=[result_1]) ev_fast = btn_fast.click(infer_fast, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], outputs=[result_1]) ev_std = btn_std.click(infer_std, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], outputs=[result_1]) ev_long = btn_long.click(infer_long, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], outputs=[result_1]) ev_max = btn_max.click(infer_max, inputs=[prompt, seed, width, height, num_inference_steps, cfg, positive_prompt, negative_prompt], outputs=[result_1]) # Trigger JS dispatcher on run button or prompt submit run_button.click(None, inputs=[width, height, num_inference_steps], outputs=[], js=js_dispatch) prompt.submit(None, inputs=[width, height, num_inference_steps], outputs=[], js=js_dispatch) # Cancel button cancels all possible events cancel_button.click(fn=None, inputs=None, outputs=None, cancels=[ev_tiny, ev_fast, ev_std, ev_long, ev_max]) if __name__ == "__main__": demo.launch()