live_preview_helpers.py

# live_preview_helpers.py

import torch
from typing import Iterator

def flux_pipe_call_that_returns_an_iterable_of_images(
    self,
    prompt: str,
    guidance_scale: float = 3.5,
    num_inference_steps: int = 28,
    width: int = 1024,
    height: int = 1024,
    generator=None,
    output_type: str = "pil",
    good_vae=None,
) -> Iterator:
    """
    Streams an iterable of images as generated by the FLUX pipeline, for use with Gradio's live preview.
    """
    # You may use your actual FLUX pipeline API if different.
    pipe = self  # usually the FLUX pipeline instance

    if generator is None:
        generator = torch.Generator(device="cpu").manual_seed(0)
    images = pipe(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        output_type=output_type,
        vae=good_vae,
    ).images

    for img in images:
        yield img