it run on cpu only
colab tpu for big ram
!pip install -U diffusers
!pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
!pip uninstall -y tensorflow && pip install tensorflow-cpu
from diffusers import FluxPipeline, FluxTransformer2DModel
import torch
Configuration
MODEL_NAME_OR_DIR = "black-forest-labs/FLUX.1-schnell"
IMAGE_OUTPUT_DIR = "/content/1"
IMAGE_PREFIX = "sign_"
DEVICE = torch.device("cpu")
If True, uses pipeline.enable_sequential_cpu_offload(). Make sure device is CPU.
USE_CPU_OFFLOAD = True
SEED = 42
IMAGE_WIDTH = 1024
IMAGE_HEIGHT = 768
NUM_STEPS = 8
NUM_IMAGES = 1
CFG = 0
PROMPT = 'dog'
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
torch_dtype=torch.float32,
use_safetensors=True, local_dir="/content/1", local_dir_use_symlinks=False,
ignore_patterns=["flux1-schnell.sft","flux1-schnell.safetensors"]).to(DEVICE)
images = pipeline(
prompt=PROMPT,
guidance_scale=CFG,
num_inference_steps=NUM_STEPS,
height=IMAGE_HEIGHT,
width=IMAGE_WIDTH,
max_sequence_length=512,
generator=torch.manual_seed(42),
num_images_per_prompt=NUM_IMAGES,
).images
from IPython.display import display
display(images[0])
black-forest-labs/FLUX.1-dev
from diffusers import FluxPipeline, FluxTransformer2DModel
import torch
Configuration
MODEL_NAME_OR_DIR = "black-forest-labs/FLUX.1-dev"
IMAGE_OUTPUT_DIR = "/content/1"
IMAGE_PREFIX = "sign_"
DEVICE = torch.device("cpu")
If True, uses pipeline.enable_sequential_cpu_offload(). Make sure device is CPU.
USE_CPU_OFFLOAD = True
SEED = 42
IMAGE_WIDTH = 1024
IMAGE_HEIGHT = 768
NUM_STEPS = 8
NUM_IMAGES = 1
CFG = 0
PROMPT = 'dog'
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
torch_dtype=torch.float32,
use_safetensors=True, local_dir="/content/1", local_dir_use_symlinks=False,
ignore_patterns=["flux1-dev.sft","flux1-dev.safetensors"]).to(DEVICE)
images = pipeline(
prompt=PROMPT,
guidance_scale=CFG,
num_inference_steps=NUM_STEPS,
height=IMAGE_HEIGHT,
width=IMAGE_WIDTH,
max_sequence_length=512,
generator=torch.manual_seed(42),
num_images_per_prompt=NUM_IMAGES,
).images
from IPython.display import display
display(images[0])