FLUX.1-schnell 6枚 バッチだと 3倍速いな

https://huggingface.co/black-forest-labs

改造して公開している120秒版
https://huggingface.co/spaces/Akjava/FLUX.1-schnell-120sec

pipe.enable_sequential_cpu_offload() 

をしないとVRAM 16GBでは動かなかった。

1024x1024 １枚づつ作成

HuggingSpaceのコードそのままだとA4000の場合

合計 16枚の画像の生成と保存にかかった時間: 639.81秒

#@spaces.GPU(duration=120)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4,guidance_scale=0):
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
            prompt = prompt, 
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generator,
            guidance_scale=guidance_scale
    ).images[0] 
    return image, seed

1024x1024 6枚づつ作成

8枚だとメモリー不足、偶数にしたかったので６枚

合計 48枚の画像の生成と保存にかかった時間: 668.82秒

#@spaces.GPU(duration=120)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4,guidance_scale=0,batch_size=1,seed_change=1000):
        
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generators = [torch.Generator(device).manual_seed(seed + j*args.seed_change) for j in range(batch_size)]
     
    output = pipe(
            prompt = prompt, 
            width = width,
            height = height,
            num_inference_steps = num_inference_steps, 
            generator = generators,
            guidance_scale=guidance_scale,num_images_per_prompt=batch_size
    )
    seed += seed_change*batch_size
    return output.images, seed

作ってみたい

試してみたら、256x256でも、そこそこいいのが出来るので、バッチでまとめて作って保存できるみたいなDemoを作ってみたい