flux pipeline restored

app.py

@@ -1,221 +1,133 @@
 # 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
+from PIL import Image
+import os # Хотя скачивание не используется, оставим на всякий случай
+# from tqdm import tqdm # Не используется в этом скрипте
 
-# Импорты из diffusers
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler, StableDiffusionPipeline
+# Импорты для FLUX ControlNet пайплайна
+from diffusers import FluxControlNetPipeline, ControlNetModel, FluxPipeline
 # 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 моделей 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"
 
-# --- Определение путей/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)
 
-# --- Загрузка моделей и создание пайплайна ---
+# --- Загрузка пайплайна FLUX ControlNet ---
 # Эта функция вызывается один раз при запуске скрипта
-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
+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}")
 
-    print(f"Загрузка базовой модели из локального файла: {base_model_path} с использованием from_single_file")
-    # Используем from_single_file для загрузки пайплайна из одного SafeTensor файла
     try:
-        pipe = StableDiffusionPipeline.from_single_file(
-            base_model_path, # Указываем путь к локальному файлу .safetensors
+        # Пайплайн 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,
-            # 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
+            # safety_checker=None # Обычно from_pretrained для FLUX пайплайна не принимает этот аргумент напрямую
         )
 
-        # Рекомендуется использовать планировщик 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 пайплайна.")
-
+        # Для FLUX планировщик специфичный, from_pretrained должен загрузить правильный.
+        print(f"Планировщик загружен: {type(pipe.scheduler).__name__}")
 
-        # Перемещаем ControlNet пайплайн на GPU, если доступно
+        # Перемещаем пайплайн на GPU, если доступно
         if torch.cuda.is_available():
-            controlnet_pipe = controlnet_pipe.to("cuda")
-            print("Финальный пайплайн перемещен на GPU.")
+            pipe = pipe.to("cuda")
+            print("Пайплайн FLUX ControlNet перемещен на GPU.")
         else:
-            print("GPU не найдено. Пайплайн будет работать на CPU (крайне медленно).") # В Space на CPU работать не будет эффективно
+            print("GPU не найдено. Пайплайн будет работать на CPU (не рекомендуется для FLUX).")
 
-        return controlnet_pipe # Возвращаем готовый пайплайн
+        print("Загрузка пайплайна FLUX ControlNet завершена успешно.")
+        return pipe # Возвращаем готовый пайплайн
 
     except Exception as e:
-        print(f"Ошибка при создании финального StableDiffusionControlNetPipeline: {e}")
-        print("Проверьте совместимость компонентов (базовая модель и ControlNet).")
-        return None # Возвращаем None, если собрать финальный пайплайн не удалось
+        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, если загрузка не удалась
 
 
-# --- Загружаем пайплайн при запуске скрипта, только если файл модели успешно скачан ---
-# Этот код выполняется после скачивания файла
-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 при ошибке
+# --- Загружаем пайплайн при запуске скрипта ---
+# Этот код выполняется один раз при старте 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.5, num_inference_steps: int = 30, controlnet_conditioning_scale: float = 1.0):
+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):
     """
-    Генерирует изображение с использованием Stable Diffusion ControlNet.
+    Генерирует изображение с использованием FLUX ControlNet.
     Принимает изображение NumPy, текст промта и другие параметры.
     Возвращает сгенерированное изображение в формате PIL Image.
     """
     # Проверяем, успешно ли загрузился пайплайн
     if pipeline is None:
          print("Попытка генерации, но пайплайн модели не загружен.")
-         return None, "Ошибка: Пайплайн модели не загружен. Проверьте логи Space."
+         return None, "Ошибка: Пайплайн модели FLUX не загружен. Проверьте логи Space и доступ к моделям."
 
     if controlnet_image is None:
         return None, "Ошибка: необходимо загрузить изображение для ControlNet."
 
-    print(f"Генерация изображения с промтом: '{prompt}'")
-    print(f"Размер входного изображения: {controlnet_image.shape}")
+    print(f"Генерация изображения FLUX с промтом: '{prompt}'")
+    print(f"Размер входного изображения для ControlNet: {controlnet_image.shape}")
 
     # Gradio возвращает изображение как numpy array. Преобразуем в PIL Image для пайплайна.
-    # diffusers ControlNet ожидают изображение в формате PIL Image или PyTorch Tensor в RGB
+    # Пайплайны 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:
-        # Здесь вы можете добавить generator=... (для сидов), width=..., height=..., etc.
-        # Передаем все параметры в вызов пайплайна
+        # Вызов пайплайна 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
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            # Другие параметры, специфичные для FLUX, могут быть доступны здесь.
+            # Проверьте сигнатуру вызова пайплайна FLUX в diffusers.
         )
 
         # Результат находится в output.images[0]
         generated_image_pil = output.images[0]
 
-        print("Генерация завершена.")
+        print("Генерация FLUX завершена.")
         return generated_image_pil, "Успех!"
     except Exception as e:
-        print(f"Ошибка при генерации: {e}")
+        print(f"Ошибка при генерации FLUX: {e}")
         # Возвращаем None и сообщение об ошибке в интерфейс Gradio
-        return None, f"Ошибка при генерации: {e}"
+        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.5, step=0.1, label="Степень соответствия промту (Guidance Scale)")
-num_inference_steps_comp = gr.Slider(minimum=10, maximum=150, value=30, step=1, label="Количество шагов (Inference Steps)")
+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="Сгенерированное изображение")
@@ -236,8 +148,8 @@ interface = gr.Interface(
         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."
+    title="FLUX ControlNet Interface on HF Space",
+    description="Загрузите изображение для ControlNet, введите промт и нажмите 'Generate'. Используются модели FLUX и FLUX ControlNet с Hugging Face Hub."
 )
 
 # Нет необходимости вызывать interface.launch() в блоке if __name__ == "__main__":