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 # Добавляем импорт tqdm

# Импорты из diffusers
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler, StableDiffusionPipeline
# from diffusers.utils import load_image # Не нужен для этого кода
# from huggingface_hub import hf_hub_download # Не нужен для этого кода

# --- Вспомогательная функция для скачивания файлов (например, с Civitai) ---
# Эта функция будет скачивать модель SafeTensor внутри Space при первом запуске
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() # Проверка на ошибки HTTP

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

        # Используем tqdm для индикатора прогресса, только если размер известен
        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 модели с Civitai
CIVITAI_SAFETENSOR_URL = "https://civitai.com/api/download/models/1413133?type=Model&format=SafeTensor&size=full&fp=fp8"
# Локальное имя файла для сохранения SafeTensor модели внутри Space
LOCAL_SAFETENSOR_FILENAME = "ultrareal_fine_tune_fp8_full.safetensors"

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

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

# --- Скачиваем SafeTensor модель (выполнится при запуске скрипта в Space) ---
print("Начинаю скачивание базовой модели...")
downloaded_base_model_path = download_file(CIVITAI_SAFETENSOR_URL, LOCAL_SAFETENSOR_FILENAME)

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

    print(f"Загрузка ControlNet модели: {controlnet_model_id}")
    # Загрузка ControlNet с Hugging Face Hub - кешируется автоматически Space
    try:
        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 # Не можем загрузить пайплайн без ControlNet

    print(f"Загрузка базовой модели из локального файла: {base_model_path} с использованием from_single_file")
    # Используем from_single_file для загрузки пайплайна из одного SafeTensor файла
    try:
        pipe = StableDiffusionPipeline.from_single_file(
            base_model_path, # Указываем путь к локальному файлу .safetensors
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            # from_single_file пытается найти конфигурацию VAE, tokenizer и scheduler.
            # Если ваша модель требует специфической конфигурации, возможно,
            # потребуется указать путь к папке с конфигом или загрузить их отдельно.
            # Для большинства Safetensor SD 1.5/2.x from_single_file работает из коробки.
        )
         # Отключение safety checker после загрузки, если он был загружен
        if hasattr(pipe, 'safety_checker') and pipe.safety_checker is not None:
             print("Отключаю safety checker...")
             pipe.safety_checker = None
             print("Safety checker отключен.")

    except Exception as e:
         print(f"Ошибка при загрузке базовой модели из файла {base_model_path}: {e}")
         print("Убедитесь, что файл не поврежден, соответствует формату StableDiffusion и from_single_file может его обработать.")
         return None # Возвращаем None, если загрузка базовой модели не удалась

    # --- Создание пайплайна StableDiffusionControlNetPipeline из компонентов ---
    # Этот блок выполняется ТОЛЬКО если базовая модель и ControlNet успешно загружены
    print("Создание финального пайплайна StableDiffusionControlNetPipeline...")
    try:
        controlnet_pipe = StableDiffusionControlNetPipeline(
            vae=pipe.vae,
            text_encoder=pipe.text_encoder,
            tokenizer=pipe.tokenizer,
            unet=pipe.unet,
            controlnet=controlnet, # Передаем загруженный ControlNet
            scheduler=pipe.scheduler, # Используем планировщик из базового пайплайна
            safety_checker=None, # Убираем safety_checker здесь при создании нового пайплайна
            feature_extractor=pipe.feature_extractor
        )

        # Рекомендуется использовать планировщик UniPC для ControlNet (или обновить существующий)
        # Обновляем планировщик в новом ControlNet пайплайне
        controlnet_pipe.scheduler = UniPCMultistepScheduler.from_config(controlnet_pipe.scheduler.config)

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


        # Перемещаем ControlNet пайплайн на GPU, если доступно
        if torch.cuda.is_available():
            controlnet_pipe = controlnet_pipe.to("cuda")
            print("Финальный пайплайн перемещен на GPU.")
        else:
            print("GPU не найдено. Пайплайн будет работать на CPU (крайне медленно).") # В Space на CPU работать не будет эффективно

        return controlnet_pipe # Возвращаем готовый пайплайн

    except Exception as e:
        print(f"Ошибка при создании финального StableDiffusionControlNetPipeline: {e}")
        print("Проверьте совместимость компонентов (базовая модель и ControlNet).")
        return None # Возвращаем None, если собрать финальный пайплайн не удалось


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


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

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

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

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

    # Выполняем рендеринг с помощью пайплайна
    try:
        # Здесь вы можете добавить generator=... (для сидов), width=..., height=..., etc.
        # Передаем все параметры в вызов пайплайна
        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
        )

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

        print("Генерация завершена.")
        return generated_image_pil, "Успех!"
    except Exception as e:
        print(f"Ошибка при генерации: {e}")
        # Возвращаем None и сообщение об ошибке в интерфейс Gradio
        return None, f"Ошибка при генерации: {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.5, step=0.1, label="Степень соответствия промту (Guidance Scale)")
num_inference_steps_comp = gr.Slider(minimum=10, maximum=150, value=30, step=1, label="Количество шагов (Inference Steps)")
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="Stable Diffusion ControlNet Interface (SafeTensor Base Model on HF Space)",
    description="Загрузите изображение для ControlNet, введите промт и нажмите 'Generate'. Используется локальная SafeTensor модель и ControlNet с Hugging Face."
)

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