FLUX.1-dev

app.py

@@ -8,26 +8,26 @@ from diffusers import FluxPipeline
 from optimization import optimize_pipeline_
 
 
-pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
+pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16).to('cuda')
 optimize_pipeline_(pipeline, "prompt")
 
 
 @spaces.GPU
-def generate_image(prompt: str):
+def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
     generator = torch.Generator(device='cuda').manual_seed(42)
     t0 = datetime.now()
-    images = []
-    for _ in range(9):
-        image = pipeline(prompt, num_inference_steps=4, generator=generator).images[0]
-        elapsed = -(t0 - (t0 := datetime.now()))
-        images += [(image, f'{elapsed.total_seconds():.2f}s')]
-        yield images
+    output = pipeline(
+        prompt=prompt,
+        num_inference_steps=28,
+        generator=generator,
+    )
+    return (output.images[0], f'{(datetime.now() - t0).total_seconds():.2f}s')
 
 
 gr.Interface(
     fn=generate_image,
     inputs=gr.Text(label="Prompt"),
-    outputs=gr.Gallery(rows=3, columns=3, height='60vh'),
+    outputs=gr.Image(),
     examples=["A cat playing with a ball of yarn"],
     cache_examples=False,
 ).launch()