Flux Fusion [4 steps]

app.py

@@ -1,86 +1,192 @@
 # app.py
 import gradio as gr
 import torch
-import numpy as np
+import requests
 from PIL import Image
-import os # Хотя скачивание не используется, оставим на всякий случай
-# from tqdm import tqdm # Не используется в этом скрипте
+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 # Не нужен для этого кода
 
-# --- Определение 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"
+# --- Вспомогательная функция для скачивания файлов ---
+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)
 
-# --- Загрузка пайплайна 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}")
+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:
-        # Пайплайн 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 пайплайна не принимает этот аргумент напрямую
-        )
+        # Загрузка 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
 
-        # Для FLUX планировщик специфичный, from_pretrained должен загрузить правильный.
-        print(f"Планировщик загружен: {type(pipe.scheduler).__name__}")
+    print(f"Попытка собрать пайплайн FLUX ControlNet, используя локальный файл: {base_model_path} как базовую модель.")
+    print("ВНИМАНИЕ: Загрузка FLUX пайплайна из одиночного SafeTensor файла методом from_single_file")
+    print("не является стандартной и может вызвать ошибки совместимости.")
 
-        # Перемещаем пайплайн на GPU, если доступно
-        if torch.cuda.is_available():
-            pipe = pipe.to("cuda")
-            print("Пайплайн FLUX ControlNet перемещен на GPU.")
-        else:
-            print("GPU не найдено. Пайплайн будет работать на CPU (не рекомендуется для FLUX).")
+    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
 
-        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, если загрузка не удалась
+         print(f"Критическая ошибка при попытке загрузить базовый FLUX пайплайн из файла {base_model_path}: {e}")
+         print("Наиболее вероятно, этот файл SafeTensor несовместим с методами загрузки FLUX в diffusers.")
+         print("Возможно, файл поврежден или не содержит ожидаемой структуры FLUX.")
+         return None
 
 
 # --- Загружаем пайплайн при запуске скрипта ---
-# Этот код выполняется один раз при старте Space
-pipeline = load_flux_pipeline(BASE_FLUX_MODEL_ID, CONTROLNET_FLUX_MODEL_ID)
+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
-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 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, "Ошибка: Пайплайн модели FLUX не загружен. Проверьте логи Space и доступ к моделям."
+         return None, "Ошибка: Пайплайн модели не загружен. Проверьте логи Space."
 
     if controlnet_image is None:
         return None, "Ошибка: необходимо загрузить изображение для ControlNet."
@@ -88,16 +194,12 @@ def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_pr
     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
+        # Проверьте документацию diffusers для FluxControlNetPipeline для точных параметров вызова
         output = pipeline(
             prompt=prompt,
             image=input_image_pil, # Входное изображение для ControlNet
@@ -105,29 +207,28 @@ def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_pr
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
-            # Другие параметры, специфичные для FLUX, могут быть доступны здесь.
-            # Проверьте сигнатуру вызова пайплайна FLUX в diffusers.
+            # Для FLUX Fusion [4 steps], количество шагов (num_inference_steps) очень низкое!
+            # Возможно, нужно использовать фиксированное значение 4, несмотря на ползунок?
         )
 
-        # Результат находится в 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 ---
-# Определяем входные и выходные элементы
-# Элементы интерфейса могут остаться теми же, так как они универсальны
+# Параметры по умолчанию подстроены под 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_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 могут отличаться
+# 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="Сгенерированное изображение")
@@ -135,8 +236,6 @@ status_text_comp = gr.Textbox(label="Статус")
 
 
 # Создаем интерфейс Gradio
-# Поскольку мы в Space, Gradio SDK сам вызовет interface.launch()
-# Нам просто нужно определить объект интерфейса
 interface = gr.Interface(
     fn=generate_image_gradio,
     inputs=[
@@ -148,9 +247,8 @@ interface = gr.Interface(
         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."
+    title="FLUX ControlNet Interface (Attempt with Civitai SafeTensor)",
+    description="Загрузите изображение для ControlNet, введите промт и нажмите 'Generate'. Попытка использовать SafeTensor 'Flux Fusion V2' с Civitai как базовую модель FLUX с ControlNet с HF."
 )
 
-# Нет необходимости вызывать interface.launch() в блоке if __name__ == "__main__":
-# Gradio SDK в Space сделает это автоматически при запуске скрипта.
+# Запуск в Space обрабатывается SDK.