trashchenkov commited on
Commit
9d6f756
·
verified ·
1 Parent(s): bfabfc6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -7
app.py CHANGED
@@ -1,14 +1,16 @@
1
  import gradio as gr
2
- import numpy as np
3
  import torch
4
- from huggingface_hub import list_models, ModelFilter
5
  from diffusers import DiffusionPipeline
6
 
7
  # Проверка доступности модели на Hugging Face Hub
8
  def is_model_available(model_id):
9
  try:
10
- models = list_models(filter=ModelFilter(model_name=model_id))
11
- return len(models) > 0
 
 
 
12
  except Exception:
13
  return False
14
 
@@ -19,8 +21,7 @@ def validate_model(model_id):
19
  if not is_model_available(model_id):
20
  raise ValueError(f"Модель '{model_id}' не найдена на Hugging Face Hub")
21
 
22
- # Дополнительные проверки модели (можно настроить)
23
- if not ("stable-diffusion" in model_id.lower() or "sdxl" in model_id.lower()):
24
  raise ValueError("Поддерживаются только Stable Diffusion и SDXL модели")
25
 
26
  # Инициализация устройства
@@ -35,6 +36,9 @@ def load_pipeline(model_id):
35
  global pipe, current_model
36
  if model_id != current_model:
37
  validate_model(model_id)
 
 
 
38
  pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
39
  pipe = pipe.to(device)
40
  current_model = model_id
@@ -96,7 +100,7 @@ css = """
96
 
97
  with gr.Blocks(css=css) as demo:
98
  with gr.Column(elem_id="col-container"):
99
- gr.Markdown("# 🎨 Text-to-Image Generator")
100
 
101
  with gr.Row():
102
  model_id = gr.Dropdown(
 
1
  import gradio as gr
 
2
  import torch
3
+ from huggingface_hub import HfApi, RepositoryNotFoundError
4
  from diffusers import DiffusionPipeline
5
 
6
  # Проверка доступности модели на Hugging Face Hub
7
  def is_model_available(model_id):
8
  try:
9
+ api = HfApi()
10
+ api.model_info(model_id)
11
+ return True
12
+ except RepositoryNotFoundError:
13
+ return False
14
  except Exception:
15
  return False
16
 
 
21
  if not is_model_available(model_id):
22
  raise ValueError(f"Модель '{model_id}' не найдена на Hugging Face Hub")
23
 
24
+ if not any(x in model_id.lower() for x in ["stable-diffusion", "sdxl"]):
 
25
  raise ValueError("Поддерживаются только Stable Diffusion и SDXL модели")
26
 
27
  # Инициализация устройства
 
36
  global pipe, current_model
37
  if model_id != current_model:
38
  validate_model(model_id)
39
+ if pipe is not None:
40
+ del pipe
41
+ torch.cuda.empty_cache()
42
  pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
43
  pipe = pipe.to(device)
44
  current_model = model_id
 
100
 
101
  with gr.Blocks(css=css) as demo:
102
  with gr.Column(elem_id="col-container"):
103
+ gr.Markdown("# 🎨 Генератор изображений по тексту")
104
 
105
  with gr.Row():
106
  model_id = gr.Dropdown(