app.py

# app.py
import gradio as gr
import torch
import numpy as np
from PIL import Image
import os # Хотя скачивание не используется, оставим на всякий случай
# from tqdm import tqdm # Не используется в этом скрипте

# Импорты для FLUX ControlNet пайплайна
from diffusers import FluxControlNetPipeline, ControlNetModel, FluxPipeline
# from diffusers.utils import load_image # Не нужен для этого кода

# --- Определение ID моделей FLUX на Hugging Face Hub ---
# Базовая модель FLUX (ОГРАНИЧЕННЫЙ ДОСТУП - требуется токен HF и доступ к репо)
BASE_FLUX_MODEL_ID = "black-forest-labs/FLUX.1-dev"
# ControlNet модель для FLUX (также на HF Hub)
CONTROLNET_FLUX_MODEL_ID = "ABDALLALSWAITI/FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8"


# Переменная для хранения пайплайна (будет загружен при запуске скрипта)
pipeline = None

# --- Загрузка пайплайна FLUX ControlNet ---
# Эта функция вызывается один раз при запуске скрипта
def load_flux_pipeline(base_model_id, controlnet_model_id):
    """Загружает пайплайн FLUX ControlNet с Hugging Face Hub."""
    print(f"Начинаю загрузку пайплайна FLUX ControlNet...")
    print(f"Базовая модель FLUX: {base_model_id}")
    print(f"ControlNet модель FLUX: {controlnet_model_id}")

    try:
        # Пайплайн FluxControlNetPipeline загружает и объединяет обе модели из репозиториев HF
        # from_pretrained автоматически использует HF_TOKEN, если он установлен как секрет в Space
        # Убедитесь, что версия diffusers поддерживает этот пайплайн и модели FLUX
        pipe = FluxControlNetPipeline.from_pretrained(
            base_model_id,
            controlnet=controlnet_model_id, # Передаем ID ControlNet модели
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            # safety_checker=None # Обычно from_pretrained для FLUX пайплайна не принимает этот аргумент напрямую
        )

        # Для FLUX планировщик специфичный, from_pretrained должен загрузить правильный.
        print(f"Планировщик загружен: {type(pipe.scheduler).__name__}")

        # Перемещаем пайплайн на GPU, если доступно
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
            print("Пайплайн FLUX ControlNet перемещен на GPU.")
        else:
            print("GPU не найдено. Пайплайн будет работать на CPU (не рекомендуется для FLUX).")

        print("Загрузка пайплайна FLUX ControlNet завершена успешно.")
        return pipe # Возвращаем готовый пайплайн

    except Exception as e:
        print(f"Ошибка при загрузке пайплайна FLUX ControlNet с Hugging Face Hub: {e}")
        print(f"Частые причины:")
        print(f"- Ваш аккаунт не имеет доступа к '{base_model_id}' (нужно зайти на страницу модели на hf.co и принять условия).")
        print(f"- Секрет HF_TOKEN неправильно установлен в настройках Space или не имеет достаточных прав.")
        print(f"- Указан неверный ID модели.")
        print(f"- Проблемы с интернет-соединением Space.")
        print(f"- Версия библиотеки diffusers слишком старая для моделей FLUX.")
        print(f"Подробности ошибки: {e}")
        return None # Возвращаем None, если загрузка не удалась


# --- Загружаем пайплайн при запуске скрипта ---
# Этот код выполняется один раз при старте Space
pipeline = load_flux_pipeline(BASE_FLUX_MODEL_ID, CONTROLNET_FLUX_MODEL_ID)


# --- Функция рендеринга для Gradio ---
# Эта функция будет вызываться интерфейсом Gradio в Space
def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_prompt: str = "", guidance_scale: float = 7.0, num_inference_steps: int = 50, controlnet_conditioning_scale: float = 1.0):
    """
    Генерирует изображение с использованием FLUX ControlNet.
    Принимает изображение NumPy, текст промта и другие параметры.
    Возвращает сгенерированное изображение в формате PIL Image.
    """
    # Проверяем, успешно ли загрузился пайплайн
    if pipeline is None:
         print("Попытка генерации, но пайплайн модели не загружен.")
         return None, "Ошибка: Пайплайн модели FLUX не загружен. Проверьте логи Space и доступ к моделям."

    if controlnet_image is None:
        return None, "Ошибка: необходимо загрузить изображение для ControlNet."

    print(f"Генерация изображения FLUX с промтом: '{prompt}'")
    print(f"Размер входного изображения для ControlNet: {controlnet_image.shape}")

    # Gradio возвращает изображение как numpy array. Преобразуем в PIL Image для пайплайна.
    # Пайплайны ControlNet обычно ожидают PIL Image в RGB.
    input_image_pil = Image.fromarray(controlnet_image).convert("RGB")

    # Выполняем рендеринг с помощью пайплайна FLUX ControlNet
    # Параметры для FLUX могут немного отличаться от SD, проверьте документацию diffusers для FluxControlNetPipeline
    # guidance_scale и num_inference_steps - стандартные параметры
    # controlnet_conditioning_scale - стандартный параметр ControlNet
    try:
        # Вызов пайплайна FLUX ControlNet
        output = pipeline(
            prompt=prompt,
            image=input_image_pil, # Входное изображение для ControlNet
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            # Другие параметры, специфичные для FLUX, могут быть доступны здесь.
            # Проверьте сигнатуру вызова пайплайна FLUX в diffusers.
        )

        # Результат находится в output.images[0]
        generated_image_pil = output.images[0]

        print("Генерация FLUX завершена.")
        return generated_image_pil, "Успех!"
    except Exception as e:
        print(f"Ошибка при генерации FLUX: {e}")
        # Возвращаем None и сообщение об ошибке в интерфейс Gradio
        return None, f"Ошибка при генерации FLUX: {e}"


# --- Настройка интерфейса Gradio ---
# Определяем входные и выходные элементы
# Элементы интерфейса могут остаться теми же, так как они универсальны
input_image_comp = gr.Image(type="numpy", label="Изображение для ControlNet (набросок, карта глубины и т.д.)")
prompt_comp = gr.Textbox(label="Промт (Prompt)")
negative_prompt_comp = gr.Textbox(label="Негативный промт (Negative Prompt)")
guidance_scale_comp = gr.Slider(minimum=1.0, maximum=20.0, value=7.0, step=0.1, label="Степень соответствия промту (Guidance Scale)")
num_inference_steps_comp = gr.Slider(minimum=10, maximum=150, value=50, step=1, label="Количество шагов (Inference Steps)") # Шаги для FLUX могут отличаться
controlnet_conditioning_scale_comp = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.05, label="Вес ControlNet (ControlNet Scale)")

output_image_comp = gr.Image(type="pil", label="Сгенерированное изображение")
status_text_comp = gr.Textbox(label="Статус")


# Создаем интерфейс Gradio
# Поскольку мы в Space, Gradio SDK сам вызовет interface.launch()
# Нам просто нужно определить объект интерфейса
interface = gr.Interface(
    fn=generate_image_gradio,
    inputs=[
        input_image_comp,
        prompt_comp,
        negative_prompt_comp,
        guidance_scale_comp,
        num_inference_steps_comp,
        controlnet_conditioning_scale_comp
    ],
    outputs=[output_image_comp, status_text_comp],
    title="FLUX ControlNet Interface on HF Space",
    description="Загрузите изображение для ControlNet, введите промт и нажмите 'Generate'. Используются модели FLUX и FLUX ControlNet с Hugging Face Hub."
)

# Нет необходимости вызывать interface.launch() в блоке if __name__ == "__main__":
# Gradio SDK в Space сделает это автоматически при запуске скрипта.