Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch | |
| import os | |
| import random | |
| import gradio as gr | |
| import torch | |
| from pipeline import generate_image, load_models | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| MAX_SEED = 2**31 - 1 | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| MODELS = load_models(device="cuda") | |
| def _duration(prompt, seed=0, randomize_seed=True, num_steps=50, *args, **kwargs): | |
| # measured on ZeroGPU: ~7.9 s cold / ~2.9 s warm for 50 steps (incl. VAE decode) | |
| return int(7 + 0.07 * int(num_steps)) | |
| def generate( | |
| prompt: str, | |
| seed: int = 0, | |
| randomize_seed: bool = True, | |
| num_steps: int = 50, | |
| cfg_scale: float = 7.0, | |
| guidance_low: float = 0.35, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Generate a 512x512 image from a text prompt with the GenFirst EiT MMDiT-XL/2 prior (variant A). | |
| Args: | |
| prompt: Text description of the image. | |
| seed: Random seed (ignored when randomize_seed is True). | |
| randomize_seed: Pick a random seed for every call. | |
| num_steps: Number of Euler sampling steps. | |
| cfg_scale: Classifier-free guidance scale. | |
| guidance_low: CFG is only applied for t in [guidance_low, 1.0] (t=1 is noise). | |
| Returns: | |
| The generated image and the seed that was used. | |
| """ | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Please enter a prompt.") | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| seed = int(seed) | |
| num_steps = int(num_steps) | |
| progress(0, desc="Sampling") | |
| image = generate_image( | |
| MODELS, prompt.strip(), seed, num_steps, float(cfg_scale), float(guidance_low), 1.0, | |
| device="cuda", autocast=True, | |
| callback=lambda i: progress((i + 1) / num_steps, desc=f"Step {i + 1}/{num_steps}"), | |
| ) | |
| return image, seed | |
| CSS = """ | |
| #col-container { max-width: 900px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| EXAMPLES = [ | |
| ["a photograph of an astronaut riding a horse"], | |
| ["a cute corgi puppy in a field of flowers, bokeh"], | |
| ["an oil painting of a lighthouse during a storm"], | |
| ["a bowl of ramen with an egg, food photography"], | |
| ["a photo of a purple wine glass and a black apple"], | |
| ["a photo of a dog right of a teddy bear"], | |
| ] | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """# GenFirst: Text-to-Image (EiT MMDiT-XL/2) | |
| 512x512 text-to-image from **GenFirst** (*Generation Before Reconstruction*): an MMDiT-XL/2 flow-matching prior with a CLIP-L text encoder, trained on an end-to-end GenFirst FLUX VAE (variant A, GenEval 0.90). | |
| [Paper](https://arxiv.org/abs/2608.29335) · [Code](https://github.com/ByteDance-Seed/GenFirst) · [Weights](https://huggingface.co/ByteDance-Seed/GenFirst) · [Project page](https://tom-zgt.github.io/GenFirst/) · weights are CC BY-NC 4.0 (non-commercial).""" | |
| ) | |
| with gr.Row(): | |
| prompt = gr.Textbox( | |
| show_label=False, placeholder="Describe the image you want…", container=False, scale=4, max_lines=3 | |
| ) | |
| run = gr.Button("Generate", variant="primary", scale=1) | |
| result = gr.Image(label="Result", type="numpy", format="png", height=512) | |
| with gr.Accordion("Advanced settings", open=False): | |
| with gr.Row(): | |
| seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| with gr.Row(): | |
| num_steps = gr.Slider(label="Sampling steps", minimum=10, maximum=100, step=1, value=50) | |
| cfg_scale = gr.Slider(label="Guidance scale (CFG)", minimum=1.0, maximum=12.0, step=0.5, value=7.0) | |
| guidance_low = gr.Slider( | |
| label="CFG interval start (t_low)", minimum=0.0, maximum=1.0, step=0.05, value=0.35, | |
| info="CFG is applied on t ∈ [t_low, 1]; the paper's GenEval setting is 0.35", | |
| ) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[prompt], | |
| outputs=[result, seed], | |
| fn=generate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| gr.on( | |
| triggers=[run.click, prompt.submit], | |
| fn=generate, | |
| inputs=[prompt, seed, randomize_seed, num_steps, cfg_scale, guidance_low], | |
| outputs=[result, seed], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |