filipstrand / mflux

A MLX port of FLUX based on the Huggingface Diffusers implementation.
656 stars 48 forks source link

Add Stream generate image functions #21

Open madroidmaq opened 3 weeks ago

madroidmaq commented 3 weeks ago

Add the stream_generate_image function to expose the results of each step of the inference for follow-up GUI interaction.

The usage would look something like this:

def stream_demo(args):
    seed = int(time.time()) if args.seed is None else args.seed

    flux = Flux1.from_alias(args.model)

    stream = flux.stream_generate_image(
        seed=seed,
        prompt=args.prompt,
        report_step=5,
        config=Config(
            num_inference_steps=args.steps,
            height=args.height,
            width=args.width,
            guidance=args.guidance,
        ),
    )

    step_count = 0

    for image in stream:
        step_count += 1

        intermediate_output_path = f"{args.output}_report_step_{step_count}.png"
        ImageUtil.save_image(image, "build/stream/" + intermediate_output_path)

    print(f"Generation complete. Total steps: {step_count}")

stream_demo()