aducsdr commited on
Commit
340204e
·
verified ·
1 Parent(s): d6fe666

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -2
app.py CHANGED
@@ -5,6 +5,38 @@ import os
5
  import numpy as np
6
  from PIL import Image
7
  from omegaconf import OmegaConf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
 
9
  from image_datasets.dataset import image_resize
10
  args = OmegaConf.load("inference_configs/inference.yaml")
@@ -19,10 +51,11 @@ def generate(image: Image.Image, edit_prompt: str):
19
 
20
  global sampler
21
  if sampler == None:
 
22
  sampler = XFluxSampler(
23
  device = device,
24
  ip_loaded=False,
25
- spatial_condition=True,
26
  clip_image_processor=None,
27
  image_encoder=None,
28
  improj=None,
@@ -202,4 +235,5 @@ def create_app():
202
  return app
203
 
204
  if __name__ == "__main__":
205
- create_app().launch(debug=False, share=False, ssr_mode=False)
 
 
5
  import numpy as np
6
  from PIL import Image
7
  from omegaconf import OmegaConf
8
+ from huggingface_hub import hf_hub_download
9
+
10
+ # --- Início: Bloco de Download Automático do Modelo ---
11
+
12
+ # Define o diretório e o caminho para os pesos do modelo
13
+ WEIGHTS_DIR = "./pretrained_weights/ByteMorpher"
14
+ MODEL_FILENAME = "dit.safetensors"
15
+ MODEL_PATH = os.path.join(WEIGHTS_DIR, MODEL_FILENAME)
16
+
17
+ # Cria o diretório se ele não existir
18
+ os.makedirs(WEIGHTS_DIR, exist_ok=True)
19
+
20
+ # Verifica se o modelo já existe antes de fazer o download
21
+ if not os.path.exists(MODEL_PATH):
22
+ print(f"Modelo não encontrado em {MODEL_PATH}. Baixando do Hugging Face Hub...")
23
+ try:
24
+ hf_hub_download(
25
+ repo_id="ByteDance-Seed/BM-Model",
26
+ filename=MODEL_FILENAME,
27
+ local_dir=WEIGHTS_DIR,
28
+ local_dir_use_symlinks=False # Recomendado para Hugging Face Spaces
29
+ )
30
+ print("Download do modelo concluído com sucesso.")
31
+ except Exception as e:
32
+ print(f"Ocorreu um erro durante o download do modelo: {e}")
33
+ # Se o download falhar, o aplicativo não poderá funcionar.
34
+ # Você pode adicionar um tratamento de erro mais robusto aqui se desejar.
35
+ else:
36
+ print(f"Modelo já existe em {MODEL_PATH}. Pulando o download.")
37
+
38
+ # --- Fim: Bloco de Download Automático do Modelo ---
39
+
40
 
41
  from image_datasets.dataset import image_resize
42
  args = OmegaConf.load("inference_configs/inference.yaml")
 
51
 
52
  global sampler
53
  if sampler == None:
54
+ # A inicialização do sampler agora ocorrerá após a confirmação de que o modelo foi baixado.
55
  sampler = XFluxSampler(
56
  device = device,
57
  ip_loaded=False,
58
+ spatial_condition=False,
59
  clip_image_processor=None,
60
  image_encoder=None,
61
  improj=None,
 
235
  return app
236
 
237
  if __name__ == "__main__":
238
+ app = create_app()
239
+ app.launch(debug=False, share=False)