stable-diffusion-xl-base-1.0_in_GoogleColab

sd-xl-base-1.0_in_GoogleColab.py

# Developed by MikyPo
# More code for DA here: https://dzen.ru/mikypo

# Установка необходимых библиотек
!pip install -q diffusers accelerate transformers pillow safetensors

# Импорт библиотеки torch
import torch

# Проверяем доступность CUDA
if torch.cuda.is_available():
    device_name = torch.cuda.get_device_name(0)
    print(f"✅ CUDA доступен: {device_name}")
    device = torch.device("cuda")
else:
    print("⚠ CUDA не доступен! Используется CPU (медленно)")
    device = torch.device("cpu")

# Импорт остальных библиотек
import matplotlib.pyplot as plt
import os
import time
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image
from google.colab import files

# Очистка памяти GPU
def clear_memory():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("🧹 Память очищена")

clear_memory()

print("\n⏬ Загрузка модели")
print("-" * 50)

# Загрузка модели SDXL
MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
print(f"Модель: {MODEL_NAME}")

start_time = time.time()
try:
    # Загрузка SDXL с оптимизацией
    pipe = DiffusionPipeline.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16,   # Обязательно FP16 для экономии памяти
        variant="fp16",              # Загружаем FP16 версию
        use_safetensors=True,        # Используем безопасный формат
        safety_checker=None,         # Убираем цензор для скорости
        requires_safety_checker=False
    )

    # Важные оптимизаци для Colab:

    # 1. CPU offload - загружает части модели в CPU
    pipe.enable_model_cpu_offload()

    # 2. Attention slicing - делит внимание на части
    pipe.enable_attention_slicing()

    # 3. Улучшенный scheduler для качества
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        algorithm_type="dpmsolver++",
        use_karras_sigmas=True
    )

    print(f"\n✅ Модель SDXL загружена за {time.time() - start_time:.1f} сек")
    print(f"Тип данных: {pipe.dtype}")
    print(f"Оптимизации: CPU offload + Attention slicing")

except torch.cuda.OutOfMemoryError:
    print("❌ Недостаточно памяти даже с оптимизациями!")
    print("Попробуйте:")
    print("1. MODEL_NAME = 'stabilityai/sdxl-turbo' (легче)")
    print("2. Перезапустите Colab и попробуйте снова")
    print("3. Используйте SD 2.1 вместо SDXL")
    raise

except Exception as e:
    print(f"❌ Ошибка загрузки модели: {e}")
    raise

# Основная функция генерации ДЛЯ SDXL
def generate_image_sdxl(params):
    """
    Генерирует одно изображение с помощью SDXL

    Параметры:
    params - словарь с параметрами генерации:
        - prompt: описание изображения (обязательно)
        - negative_prompt: что исключить (опционально)
        - steps: шагов генерации (20-40 для SDXL)
        - guidance_scale: сила следования промпту (3-7 для SDXL)
        - width, height: размер изображения (рекомендуется 1024x1024)
        - seed: для воспроизводимости (опционально)
        - save_to_drive: сохранять в Google Drive (True/False)
    """

    # Извлекаем параметры с значениями по умолчанию ДЛЯ SDXL
    prompt = params.get('prompt', '')
    if not prompt:
        raise ValueError("Не указан промпт (prompt)")

    negative_prompt = params.get('negative_prompt', '')
    steps = params.get('steps', 30)
    guidance_scale = params.get('guidance_scale', 5.0)  # SDXL: 3-7 оптимально
    width = params.get('width', 1024)  # SDXL родной размер
    height = params.get('height', 1024)
    seed = params.get('seed', None)
    save_to_drive = params.get('save_to_drive', False)

    print()
    print(f"\n🎨 Генерация изображения SDXL")
    print(f"Промпт: {prompt[:100]}...")
    print(f"Размер: {width}x{height} px")
    print(f"Шаги: {steps}, Guidance scale: {guidance_scale}")

    # Проверка параметров для SDXL
    if guidance_scale > 7.5:
        print(f"⚠ Внимание: guidance_scale={guidance_scale} высоковат для SDXL")
        print("  Рекомендуется 3-7")

    if steps > 50:
        print(f"⚠ Внимание: steps={steps} много для SDXL")
        print("  Рекомендуется 20-40")

    if negative_prompt:
        print(f"Отрицательный промпт: {negative_prompt[:80]}...")

    # Устанавливаем seed если указан
    generator = None
    if seed is not None:
        # Для SDXL с CPU offload используем CPU генератор
        generator = torch.Generator("cpu").manual_seed(seed)
        print(f"Seed: {seed}")

    # Генерация изображения
    gen_start = time.time()

    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            generator=generator,
            num_images_per_prompt=1,
            output_type="pil"  # Явно указываем тип вывода
        )

        image = result.images[0]
        gen_time = time.time() - gen_start
        print(f"\n✅ Изображение сгенерировано за {gen_time:.1f} сек")

    except torch.cuda.OutOfMemoryError:
        print("\n❌ Недостаточно памяти GPU!")
        print("Попробуйте:")
        print("  - Уменьшить размер до 768x768")
        print("  - Уменьшить количество шагов до 20-25")
        print("  - Перезапустить среду выполнения")
        return None

    # Показываем изображение
    print()
    print("\n🖼️ Результат:")
    plt.figure(figsize=(12, 12))
    plt.imshow(image)
    plt.axis('off')
    plt.show()

    # Сохраняем файл
    timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
    filename = f"sdxl_{width}x{height}px_{timestamp}.png"

    # Создаем папку для результатов
    os.makedirs("/content/sdxl_results", exist_ok=True)
    filepath = f"/content/sdxl_results/{filename}"
    image.save(filepath, quality=95)
    print(f"💾 Файл сохранен: {filename}")

    # Сохранение в Google Drive
    if save_to_drive:
        try:
            from google.colab import drive
            drive.mount('/content/drive')

            drive_path = f"/content/drive/MyDrive/sdxl_generations"
            os.makedirs(drive_path, exist_ok=True)

            drive_filepath = f"{drive_path}/{filename}"
            image.save(drive_filepath, quality=95)
            print(f"📁 Также сохранено в Google Drive: {drive_filepath}")

        except Exception as e:
            print(f"⚠ Не удалось сохранить в Google Drive: {e}")

    # Предлагаем скачать
    download = input("\nСкачать изображение на компьютер? (y/n): ").lower()
    if download == 'y':
        files.download(filepath)
        print("✅ Изображение скачано!")

    # Очистка памяти
    clear_memory()

# ============================================
# ПАРАМЕТРЫ ДЛЯ ГЕНЕРАЦИИ
# ============================================

params = {
    'prompt': 'Surrealism painting Warlock on marketplace',
    'negative_prompt': 'realism, photo, realistic, 3d, blurry, low quality, deformed, ugly',
    'steps': 40,         # SDXL: 25-40 оптимально
    'guidance_scale': 7, # SDXL: 3-7 оптимально
    'width': 1024,
    'height': 1024,
    'seed': None,
    'save_to_drive': False
}

# ============================================
# ИНФОРМАЦИЯ И ЗАПУСК
# ============================================

print("🚀 SDXL Готова к генерации")
print("-"*60)
print("Параметры:")
for key, value in params.items():
    print(f"  {key}: {value}")

print("\n" + "-"*60)
print("⚠ Важно для SDXL:")
print("  • Guidance scale 3-7 (не 7-12 как в SD 1.5)")
print("  • Steps 20-40 достаточно для качества")
print("  • Размер 1024x1024 оптимальный")
print("  • Более детальные промпты дают лучший результат")
print("-"*60)

print("\nНачинаю генерацию...")

# Запуск генерации
generate_image_sdxl(params)

print("\n" + "-"*60)
print("✅ Генерация завершена")