fix path model

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import torch
 import requests
 from PIL import Image
-from io import BytesIO
 import numpy as np
 import os
 from tqdm import tqdm # Добавляем импорт tqdm
@@ -30,11 +29,22 @@ def download_file(url, local_filename):
         total_size_in_bytes = int(response.headers.get('content-length', 0))
         block_size = 8192 # 8 Kibibytes
 
-        with tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True, desc=f"Скачивание {local_filename}") as progress_bar:
-            with open(local_filename, 'wb') as f:
-                for chunk in response.iter_content(chunk_size=block_size):
+        # Используем 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)
+                f.write(chunk)
+
+        if progress_bar:
+            progress_bar.close()
 
         print(f"Скачивание завершено: {local_filename}")
         return local_filename
@@ -55,36 +65,41 @@ 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)
 
-if not downloaded_base_model_path or not os.path.exists(downloaded_base_model_path):
-    # Если скачивание не удалось или файл не существует после попытки
-    print(f"Критическая ошибка: Не удалось получить файл базовой модели по пути: {LOCAL_SAFETENSOR_FILENAME}")
-    print("Проверьте логи Space на наличие ошибок скачивания.")
-    # Возможно, здесь стоит выбросить исключение или как-то иначе остановить приложение
-    # Для примера, просто присвоим None и приложение не сможет загрузить пайплайн
-    pipeline = None
-else:
-    # --- Загрузка моделей и создание пайплайна ---
-    def load_pipeline_components(base_model_path, controlnet_model_id):
-        """Загружает базовую модель из локального файла, ControlNet и собирает пайплайн."""
-        print(f"Загрузка ControlNet модели: {controlnet_model_id}")
-        # Загрузка ControlNet с Hugging Face Hub - кешируется автоматически Space
+# --- Загрузка моделей и создание пайплайна ---
+# Эта функция вызывается один раз при запуске скрипта
+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")
+    print(f"Загрузка базовой модели из локального файла: {base_model_path} с использованием from_single_file")
     # Используем from_single_file для загрузки пайплайна из одного SafeTensor файла
-    # Этот метод специально предназначен для загрузки чекпойнтов (Safetensors/CKPT)
     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 часто пытается загрузить необходимую конфигурацию
-            # (tokenizer, scheduler, VAE) из стандартных источников или default config.
-            # Если у вас очень нестандартная модель, может потребоваться указать config_files=...
-            # Но для большинства SD 1.5/2.x моделей from_single_file работает хорошо.
+            # 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:
@@ -94,14 +109,13 @@ else:
 
     except Exception as e:
          print(f"Ошибка при загрузке базовой модели из файла {base_model_path}: {e}")
-         print("Убедитесь, что файл не поврежден и соответствует формату StableDiffusion.")
-         # Важно: Если загрузка базовой модели не удалась, мы не можем продолжить.
-         # Возвращаем None или выбрасываем исключение.
-         return None # Возвращаем None, чтобы функция generate_image_gradio могла проверить это.
-
-        # Теперь объединяем базовый пайплайн с ControlNet
-        # Создаем StableDiffusionControlNetPipeline на основе загруженного базового пайплайна
-        print("Создание пайплайна StableDiffusionControlNetPipeline...")
+         print("Убедитесь, что файл не поврежден, соответствует формату StableDiffusion и from_single_file может его обработать.")
+         return None # Возвращаем None, если загрузка базовой модели не удалась
+
+    # --- Создание пайплайна StableDiffusionControlNetPipeline из компонентов ---
+    # Этот блок выполняется ТОЛЬКО если базовая модель и ControlNet успешно загружены
+    print("Создание финального пайплайна StableDiffusionControlNetPipeline...")
+    try:
         controlnet_pipe = StableDiffusionControlNetPipeline(
             vae=pipe.vae,
             text_encoder=pipe.text_encoder,
@@ -109,7 +123,7 @@ else:
             unet=pipe.unet,
             controlnet=controlnet, # Передаем загруженный ControlNet
             scheduler=pipe.scheduler, # Используем планировщик из базового пайплайна
-            safety_checker=None,
+            safety_checker=None, # Убираем safety_checker здесь при создании нового пайплайна
             feature_extractor=pipe.feature_extractor
         )
 
@@ -117,22 +131,35 @@ else:
         # Обновляем планировщик в новом 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.")
+            print("Финальный пайплайн перемещен на GPU.")
         else:
-            print("GPU не найдено. Пайплайн будет работать на CPU (будет медленно).") # В Space на CPU работать не будет эффективно
+            print("GPU не найдено. Пайплайн будет работать на CPU (крайне медленно).") # В Space на CPU работать не будет эффективно
 
-        return controlnet_pipe
+        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 ---
@@ -145,6 +172,7 @@ def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_pr
     """
     # Проверяем, успешно ли загрузился пайплайн
     if pipeline is None:
+         print("Попытка генерации, но пайплайн модели не загружен.")
          return None, "Ошибка: Пайплайн модели не загружен. Проверьте логи Space."
 
     if controlnet_image is None:
@@ -154,10 +182,13 @@ def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_pr
     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
@@ -165,7 +196,6 @@ 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
-            # Здесь можно добавить generator=... (для сидов), width=..., height=..., etc.
         )
 
         # Результат находится в output.images[0]
@@ -175,6 +205,7 @@ def generate_image_gradio(controlnet_image: np.ndarray, prompt: str, negative_pr
         return generated_image_pil, "Успех!"
     except Exception as e:
         print(f"Ошибка при генерации: {e}")
+        # Возвращаем None и сообщение об ошибке в интерфейс Gradio
         return None, f"Ошибка при генерации: {e}"
 
 
@@ -193,7 +224,7 @@ status_text_comp = gr.Textbox(label="Статус")
 
 # Создаем интерфейс Gradio
 # Поскольку мы в Space, Gradio SDK сам вызовет interface.launch()
-# Нам просто нужно определить интерфейс
+# Нам просто нужно определить объект интерфейса
 interface = gr.Interface(
     fn=generate_image_gradio,
     inputs=[
@@ -205,13 +236,9 @@ interface = gr.Interface(
         controlnet_conditioning_scale_comp
     ],
     outputs=[output_image_comp, status_text_comp],
-    title="Stable Diffusion ControlNet Interface (SafeTensor Base Model)",
+    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 сделает это автоматически.
-# Если вы оставите if __name__ == "__main__": interface.launch(), оно тоже будет работать,
-# но в среде Space это менее критично, чем при локальном запуске.
-# Для ясности в Space можно убрать блок if __name__ == "__main__":
-# Я оставил его в коде выше, но знайте, что SDK вызовет interface.launch() независимо от него.
+# Нет необходимости вызывать interface.launch() в блоке if __name__ == "__main__":
+# Gradio SDK в Space сделает это автоматически при запуске скрипта.