Запускаем FLUX 1 Dev в Google Colab28.03.2025 11:46

Ранее я уже делал статью про запуск в пару кликов моделей Stable Diffusion в Google Colab с помощью Fooocus (способ актуальный), сегодня мы проделаем похожее с моделью FLUX 1 Dev, но без web интерфейса.

Шаг 1
Заходим в Google Colab, создаем новый блокнот и меняем среду выполнения на «Графический процессор Т4», подключаемся к ней.

Меняем среду выполнения
Меняем среду выполнения
Подключаемся к среде выполнения (верхний правый угол экрана)
Подключаемся к среде выполнения (верхний правый угол экрана)

Теперь добавляем первый блок кода, всего их будет два.

%cd /content
!git clone -b totoro3 https://github.com/camenduru/ComfyUI /content/TotoroUI
%cd /content/TotoroUI

!pip install -q torchsde==0.2.6 einops diffusers accelerate xformers==0.0.29.post3
!apt -y install -qq aria2

!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/flux1-dev-fp8.safetensors -d /content/TotoroUI/models/unet -o flux1-dev-fp8.safetensors
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/ae.sft -d /content/TotoroUI/models/vae -o ae.sft
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/clip_l.safetensors -d /content/TotoroUI/models/clip -o clip_l.safetensors
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/t5xxl_fp8_e4m3fn.safetensors -d /content/TotoroUI/models/clip -o t5xxl_fp8_e4m3fn.safetensors

import random
import torch
import numpy as np
from PIL import Image
import nodes
from nodes import NODE_CLASS_MAPPINGS
from totoro_extras import nodes_custom_sampler
from totoro import model_management

DualCLIPLoader = NODE_CLASS_MAPPINGS["DualCLIPLoader"]()
UNETLoader = NODE_CLASS_MAPPINGS["UNETLoader"]()
RandomNoise = nodes_custom_sampler.NODE_CLASS_MAPPINGS["RandomNoise"]()
BasicGuider = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicGuider"]()
KSamplerSelect = nodes_custom_sampler.NODE_CLASS_MAPPINGS["KSamplerSelect"]()
BasicScheduler = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicScheduler"]()
SamplerCustomAdvanced = nodes_custom_sampler.NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
EmptyLatentImage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()

with torch.inference_mode():
    clip = DualCLIPLoader.load_clip("t5xxl_fp8_e4m3fn.safetensors", "clip_l.safetensors", "flux")[0]
    unet = UNETLoader.load_unet("flux1-dev-fp8.safetensors", "fp8_e4m3fn")[0]
    vae = VAELoader.load_vae("ae.sft")[0]

def closestNumber(n, m):
    q = int(n / m)
    n1 = m * q
    if (n * m) > 0:
        n2 = m * (q + 1)
    else:
        n2 = m * (q - 1)
    if abs(n - n1) < abs(n - n2):
        return n1
    return n2

Жмем на кнопку запуска

Запуск кода
Запуск кода

Начнется процесс загрузки моделей и установки необходимых пакетов, это займет около 5 минут.

Шаг 2

Копируем отсюда и добавляем второй блок кода.

sdelatraz = 1

while sdelatraz < 88:
  with torch.inference_mode():
      positive_prompt = "A brown bear carries a toilet bowl in his hands and smiles"
      width = 1024
      height = 1024
      seed = 0
      steps = 20
      sampler_name = "euler"
      scheduler = "simple"

      if seed == 0:
          seed = random.randint(0, 18446744073709551615)
      print(seed)

      soseed = str(seed)
      sdelatraz += 1

      cond, pooled = clip.encode_from_tokens(clip.tokenize(positive_prompt), return_pooled=True)
      cond = [[cond, {"pooled_output": pooled}]]
      noise = RandomNoise.get_noise(seed)[0]
      guider = BasicGuider.get_guider(unet, cond)[0]
      sampler = KSamplerSelect.get_sampler(sampler_name)[0]
      sigmas = BasicScheduler.get_sigmas(unet, scheduler, steps, 1.0)[0]
      latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0]
      sample, sample_denoised = SamplerCustomAdvanced.sample(noise, guider, sampler, sigmas, latent_image)
      model_management.soft_empty_cache()
      decoded = VAEDecode.decode(vae, sample)[0].detach()
      Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0]).save("/content/" + soseed + "flux.png")

  Image.fromarray(np.array(decoded*255, dtype=np.uint8)[0])
print("Дело сделано")

В пятой строчке кода positive_prompt = «Сюда нужно написать свой промт» , далее запускаем этот блок кода. Начнется генерация изображений с случайным seed, изображения будут сохранятся в разделе файлы, оттуда их можно скачать. Количество изображений необходимых для генерации можно изменить в третьей строчке кода.

Готовые изображения
Готовые изображения

Есть возможность автоматически сохранять готовые изображения сразу на свой Google Диск, для этого вам придется добавить немного кода, но здесь справитесь и без меня я думаю.

Генерация изображений
Генерация изображений

Дальше я буду искать способы запуска квантованных моделей FLUX 1 Dev в Google Colab, они работают значительно быстрее. Например на моей RTX 3060 12 gb создание одного изображения с помощью квантованной версии FLUX 1 Dev 8fp в ComfyUI занимает почти в 4 раза меньше времени (семплер, количество шагов и размер изображения идентичны), при этом нагрев видеокарты ограничен до 60 градусов (она работает на процентов 60–70% от максимальной мощности). Возможно дело в ComfyUI, будем проверять.

© Habrahabr.ru