app.py

# app.py
import gradio as gr
import torch
import requests
from PIL import Image
import numpy as np
import os
from tqdm import tqdm

# Импорты для FLUX ControlNet пайплайна
# Возможно, потребуются дополнительные импорты компонентов FLUX, если from_single_file не сработает
from diffusers import FluxControlNetPipeline, ControlNetModel, FluxPipeline
# from diffusers.utils import load_image # Не нужен для этого кода

# --- Вспомогательная функция для скачивания файлов ---
def download_file(url, local_filename):
    """Скачивает файл по URL с индикатором прогресса."""
    print(f"Скачиваю {url} в {local_filename}...")
    if os.path.exists(local_filename):
        print(f"Файл уже существует: {local_filename}. Пропускаю скачивание.")
        return local_filename

    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()

        total_size_in_bytes = int(response.headers.get('content-length', 0))
        block_size = 8192

        if total_size_in_bytes > 0:
             progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True, desc=f"Скачивание {local_filename}")
        else:
             print("Размер файла неизвестен, скачивание без индикатора прогресса.")
             progress_bar = None

        with open(local_filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=block_size):
                if progress_bar:
                    progress_bar.update(len(chunk))
                f.write(chunk)

        if progress_bar:
            progress_bar.close()

        print(f"Скачивание завершено: {local_filename}")
        return local_filename
    except requests.exceptions.RequestException as e:
        print(f"Ошибка скачивания {url}: {e}")
        return None
    except Exception as e:
        print(f"Произошла другая ошибка при скачивании: {e}")
        return None

# --- Определение путей/ID моделей ---
# URL SafeTensor модели "Flux Fusion V2" с Civitai (FP8)
CIVITAI_FLUX_FUSION_URL = "https://civitai.com/api/download/models/936565?type=Model&format=SafeTensor&fp=fp8"
# Локальное имя файла для сохранения SafeTensor модели
LOCAL_FLUX_FUSION_FILENAME = "flux_fusion_v2_fp8.safetensors"

# ControlNet модель для FLUX с Hugging Face
CONTROLNET_FLUX_MODEL_ID = "ABDALLALSWAITI/FLUX.1-dev-ControlNet-Union-Pro-2.0-fp8"

# Переменная для хранения пайплайна
pipeline = None
downloaded_base_model_path = None

# --- Скачиваем SafeTensor модель с Civitai ---
print("Начинаю скачивание базовой модели с Civitai...")
downloaded_base_model_path = download_file(CIVITAI_FLUX_FUSION_URL, LOCAL_FLUX_FUSION_FILENAME)

# --- Загрузка моделей и создание пайплайна ---
# Эта функция вызывается один раз при запуске скрипта
def load_pipeline_components(base_model_path, controlnet_model_id):
    """
    Загружает ControlNet с HF и пытается собрать пайплайн FLUX,
    используя локальный SafeTensor как базовую модель.
    """
    if not base_model_path or not os.path.exists(base_model_path):
        print(f"Ошибка загрузки: Файл базовой модели не найден по пути: {base_model_path}")
        return None

    print(f"Загрузка ControlNet модели FLUX с HF Hub: {controlnet_model_id}")
    try:
        # Загрузка ControlNet для FLUX
        controlnet = ControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
    except Exception as e:
        print(f"Ошибка загрузки ControlNet модели с HF Hub: {controlnet_model_id}. Проверьте ID или соединение.")
        print(f"Ошибка: {e}")
        return None

    print(f"Попытка собрать пайплайн FLUX ControlNet, используя локальный файл: {base_model_path} как базовую модель.")
    print("ВНИМАНИЕ: Загрузка FLUX пайплайна из одиночного SafeTensor файла методом from_single_file")
    print("не является стандартной и может вызвать ошибки совместимости.")

    try:
        # !!! ЭТО САМАЯ ПРОБЛЕМНАЯ ЧАСТЬ !!!
        # from_single_file разработан для SD. Попытка использовать его для FLUX SafeTensor может не сработать.
        # from_pretrained для FluxControlNetPipeline ожидает ID репозитория HF или локальную ПАПКУ.
        # Здесь мы пытаемся передать локальный *файл*. Это нестандартно.
        # Возможно, придется явно указывать тип модели или компоненты, если from_single_file не сработает.
        # Например: FluxPipeline.from_single_file() если такой метод есть и работает для FLUX.
        # Или даже собрать вручную: FluxPipeline(transformer=..., vae=..., ...).from_single_file(...)
        # Попробуем передать файл в from_pretrained, хотя он обычно ждет папку/ID.
        # Или попытаемся использовать from_single_file, хотя он для SD.
        # Основываясь на предыдущем опыте, from_single_file "пытается" понять структуру.
        # Давайте попробуем from_single_file, но с большим сомнением в успехе для FLUX.

        # Попытка 1: from_single_file (наиболее вероятный источник ошибок для FLUX SafeTensor)
        # УКАЗЫВАЕМ ЯВНО controlnet=None при загрузке БАЗОВОГО пайплайна из файла
        # ControlNetModel передадим позже при создании FluxControlNetPipeline
        base_pipe = FluxPipeline.from_single_file(
             base_model_path,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             # Возможно, придется передавать явно другие компоненты, если они не в файле
             # controlnet=None # from_single_file не принимает controlnet
        )
        print("Успешно загружен базовый FLUX пайплайн из SafeTensor файла методом from_single_file (если это сообщение видно).")


        # --- Создание финального пайплайна FluxControlNetPipeline из компонентов ---
        # Собираем пайплайн, используя компоненты из базового пайплайна и ControlNet
        print("Собираю финальный FluxControlNetPipeline...")
        # Нужно убедиться, что у base_pipe есть все необходимые для FLUX компоненты (transformer, vae, etc.)
        # from_single_file мог загрузить только часть
        try:
            controlnet_pipe = FluxControlNetPipeline(
                transformer=base_pipe.transformer, # Основной компонент FLUX
                vae=base_pipe.vae,
                text_encoder=base_pipe.text_encoder, # У FLUX есть text_encoder, но другой, не как у SD CLIP
                tokenizer=base_pipe.tokenizer,
                scheduler=base_pipe.scheduler,
                controlnet=controlnet, # Передаем загруженный FLUX ControlNet
                feature_extractor=base_pipe.feature_extractor if hasattr(base_pipe, 'feature_extractor') else None, # Копируем feature_extractor
                image_processor=base_pipe.image_processor if hasattr(base_pipe, 'image_processor') else None, # Копируем image_processor
            )

            # Планировщик должен быть FLUX-совместимым, from_single_file или from_pretrained должны его загрузить.
            print(f"Финальный планировщик: {type(controlnet_pipe.scheduler).__name__}")

            # Удаляем старый объект пайплайна для освобождения памяти GPU
            del base_pipe
            if torch.cuda.is_available():
                 torch.cuda.empty_cache()
                 print("Память GPU очищена после создания ControlNet пайплайна.")

            # Перемещаем пайплайн на GPU
            if torch.cuda.is_available():
                controlnet_pipe = controlnet_pipe.to("cuda")
                print("Финальный пайплайн FLUX ControlNet перемещен на GPU.")
            else:
                 print("GPU не найдено. Пайплайн на CPU.")

            print("Сборка финального пайплайна FLUX ControlNet завершена успешно.")
            return controlnet_pipe

        except Exception as e:
            print(f"Ошибка при сборке финального FluxControlNetPipeline: {e}")
            print("Проверьте, что базовая модель, загруженная из SafeTensor, содержит все компоненты FLUX (transformer, vae, text_encoder и т.д.).")
            print("Возможно, from_single_file не смог загрузить все необходимые компоненты FLUX из этого файла.")
            return None


    except Exception as e:
         print(f"Критическая ошибка при попытке загрузить базовый FLUX пайплайн из файла {base_model_path}: {e}")
         print("Наиболее вероятно, этот файл SafeTensor несовместим с методами загрузки FLUX в diffusers.")
         print("Возможно, файл поврежден или не содержит ожидаемой структуры FLUX.")
         return None


# --- Загружаем пайплайн при запуске скрипта ---
if downloaded_base_model_path and os.path.exists(downloaded_base_model_path):
    pipeline = load_pipeline_components(downloaded_base_model_path, CONTROLNET_FLUX_MODEL_ID)
else:
    print("Пропуск загрузки пайплайна из-за ошибки скачивания или отсутствия файла.")
    pipeline = None


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

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

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

    input_image_pil = Image.fromarray(controlnet_image).convert("RGB")

    # Выполняем рендеринг с помощью пайплайна FLUX ControlNet
    try:
        # Вызов пайплайна FLUX ControlNet
        # Проверьте документацию diffusers для FluxControlNetPipeline для точных параметров вызова
        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 Fusion [4 steps], количество шагов (num_inference_steps) очень низкое!
            # Возможно, нужно использовать фиксированное значение 4, несмотря на ползунок?
        )

        generated_image_pil = output.images[0]

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


# --- Настройка интерфейса Gradio ---
# Параметры по умолчанию подстроены под Flux Fusion [4 steps]
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 для FLUX Fusion может быть ниже, чем для SD
guidance_scale_comp = gr.Slider(minimum=0.0, maximum=10.0, value=5.0, step=0.1, label="Степень соответствия промту (Guidance Scale)")
# Количество шагов для FLUX Fusion [4 steps] ОЧЕНЬ низкое
num_inference_steps_comp = gr.Slider(minimum=1, maximum=20, value=4, step=1, label="Количество шагов (Inference Steps) [для FLUX Fusion V2 обычно 4]")
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
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 (Attempt with Civitai SafeTensor)",
    description="Загрузите изображение для ControlNet, введите промт и нажмите 'Generate'. Попытка использовать SafeTensor 'Flux Fusion V2' с Civitai как базовую модель FLUX с ControlNet с HF."
)

# Запуск в Space обрабатывается SDK.