zerogpu-aoti/FLUX.1-dev-base

1from datetime import datetime2 3import gradio as gr4import spaces5import torch6from diffusers import FluxPipeline7 8from optimization import optimize_pipeline_9 10 11pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16).to('cuda')12# optimize_pipeline_(pipeline, "prompt")13 14 15@spaces.GPU16def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):17    generator = torch.Generator(device='cuda').manual_seed(42)18    t0 = datetime.now()19    output = pipeline(20        prompt=prompt,21        num_inference_steps=28,22        generator=generator,23    )24    return [(output.images[0], f'{(datetime.now() - t0).total_seconds():.2f}s')]25 26 27gr.Interface(28    fn=generate_image,29    inputs=gr.Text(label="Prompt"),30    outputs=gr.Gallery(),31    examples=["A cat playing with a ball of yarn"],32    cache_examples=False,33).launch()34