Superigni commited on
Commit
136d95f
·
verified ·
1 Parent(s): 0adeb3f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +22 -6
app.py CHANGED
@@ -74,14 +74,30 @@ else:
74
  # Загрузка ControlNet с Hugging Face Hub - кешируется автоматически Space
75
  controlnet = ControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
76
 
77
- print(f"Загрузка базовой модели из локального файла: {base_model_path}")
78
- # Загружаем базовую модель из локального SafeTensor файла
79
- # diffusers умеет загружать локальные файлы .safetensors
80
- pipe = StableDiffusionPipeline.from_pretrained(
81
- base_model_path, # Указываем путь к локальному файлу внутри Space
 
82
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
83
- safety_checker=None # Отключение safety checker для скорости (используйте с осторожностью!)
 
 
 
84
  )
 
 
 
 
 
 
 
 
 
 
 
 
85
 
86
  # Теперь объединяем базовый пайплайн с ControlNet
87
  # Создаем StableDiffusionControlNetPipeline на основе загруженного базового пайплайна
 
74
  # Загрузка ControlNet с Hugging Face Hub - кешируется автоматически Space
75
  controlnet = ControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
76
 
77
+ print(f"Загрузка базовой модели из локального файла: {base_model_path} с использованием from_single_file")
78
+ # Используем from_single_file для загрузки пайплайна из одного SafeTensor файла
79
+ # Этот метод специально предназначен для загрузки чекпойнтов (Safetensors/CKPT)
80
+ try:
81
+ pipe = StableDiffusionPipeline.from_single_file(
82
+ base_model_path, # Указываем путь к локальному файлу .safetensors
83
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
84
+ # from_single_file часто пытается загрузить необходимую конфигурацию
85
+ # (tokenizer, scheduler, VAE) из стандартных источников или default config.
86
+ # Если у вас очень нестандартная модель, может потребоваться указать config_files=...
87
+ # Но для большинства SD 1.5/2.x моделей from_single_file работает хорошо.
88
  )
89
+ # Отключение safety checker после загрузки, если он был загружен
90
+ if hasattr(pipe, 'safety_checker') and pipe.safety_checker is not None:
91
+ print("Отключаю safety checker...")
92
+ pipe.safety_checker = None
93
+ print("Safety checker отключен.")
94
+
95
+ except Exception as e:
96
+ print(f"Ошибка при загрузке базовой модели из файла {base_model_path}: {e}")
97
+ print("Убедитесь, что файл не поврежден и соответствует формату StableDiffusion.")
98
+ # Важно: Если загрузка базовой модели не удалась, мы не можем продолжить.
99
+ # Возвращаем None или выбрасываем исключение.
100
+ return None # Возвращаем None, чтобы функция generate_image_gradio могла проверить это.
101
 
102
  # Теперь объединяем базовый пайплайн с ControlNet
103
  # Создаем StableDiffusionControlNetPipeline на основе загруженного базового пайплайна