?????????????????

code python

how to use flux1-dev-bnb-nf4 with python?

from diffusers import FluxPipeline
import torch

ckpt_id = "black-forest-labs/FLUX.1-schnell"
prompt = [
"an astronaut riding a horse",
# more prompts here
]
height, width = 1024, 1024

denoising

pipe = FluxPipeline.from_pretrained(
ckpt_id,
torch_dtype=torch.bfloat16,
)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)

image = pipe(
prompt,
num_inference_steps=1,
guidance_scale=0.0,
height=height,
width=width,
).images[0]
print('Max mem allocated (GB) while denoising:', torch.cuda.max_memory_allocated() / (1024 ** 3))

import matplotlib.pyplot as plt
plt.imshow(image)
plt.show()

how to use flux1-dev-bnb-nf4
instead
black-forest-labs/FLUX.1-schnell

؟؟؟؟؟؟؟؟؟؟؟؟؟؟؟؟؟؟

??????

import os
import gc
import torch
from diffusers import FluxPipeline
import matplotlib.pyplot as plt
gc.collect()
torch.cuda.empty_cache()

Define model and prompt

ckpt_id = "magespace/FLUX.1-dev-bnb-nf4"
prompt = ["an astronaut riding a horse"]
height, width = 64, 64
gc.collect()
torch.cuda.empty_cache()

Define the path to your folder

model_path = "/content/a" # Replace with your desired path
gc.collect()
torch.cuda.empty_cache()

Create the folder if it doesn't exist

os.makedirs(model_path, exist_ok=True)
gc.collect()
torch.cuda.empty_cache()

Load the pipeline, specifying the cache directory

pipe = FluxPipeline.from_pretrained(
ckpt_id,
torch_dtype=torch.bfloat16,
cache_dir=model_path,
)
gc.collect()
torch.cuda.empty_cache()

Enable memory optimization techniques

pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.enable_sequential_cpu_offload()
gc.collect()
torch.cuda.empty_cache()

Perform inference (denoising)

image = pipe(
prompt,
num_inference_steps=1,
guidance_scale=0.0,
height=height,
width=width,
).images[0]
gc.collect()
torch.cuda.empty_cache()

Print memory usage

print('Max mem allocated (GB) while denoising:', torch.cuda.max_memory_allocated() / (1024 ** 3))
gc.collect()
torch.cuda.empty_cache()

Display the image

plt.imshow(image)
plt.show()

Your session crashed after using all available RAM.

colab t4