it run on cpu only

#200
by sdyy - opened

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])

Sign up or log in to comment