Update app.py

app.py

@@ -7,17 +7,12 @@ from PIL import Image
 from omegaconf import OmegaConf
 from huggingface_hub import hf_hub_download
 
-# --- Início: Bloco de Download Automático do Modelo ---
-
-# Define o diretório e o caminho para os pesos do modelo
+# --- Bloco de Download Automático do Modelo ---
 WEIGHTS_DIR = "./pretrained_weights/ByteMorpher"
 MODEL_FILENAME = "dit.safetensors"
 MODEL_PATH = os.path.join(WEIGHTS_DIR, MODEL_FILENAME)
-
-# Cria o diretório se ele não existir
 os.makedirs(WEIGHTS_DIR, exist_ok=True)
 
-# Verifica se o modelo já existe antes de fazer o download
 if not os.path.exists(MODEL_PATH):
     print(f"Modelo não encontrado em {MODEL_PATH}. Baixando do Hugging Face Hub...")
     try:
@@ -25,17 +20,14 @@ if not os.path.exists(MODEL_PATH):
             repo_id="ByteDance-Seed/BM-Model",
             filename=MODEL_FILENAME,
             local_dir=WEIGHTS_DIR,
-            local_dir_use_symlinks=False  # Recomendado para Hugging Face Spaces
+            local_dir_use_symlinks=False
         )
         print("Download do modelo concluído com sucesso.")
     except Exception as e:
         print(f"Ocorreu um erro durante o download do modelo: {e}")
-        # Se o download falhar, o aplicativo não poderá funcionar.
-        # Você pode adicionar um tratamento de erro mais robusto aqui se desejar.
 else:
     print(f"Modelo já existe em {MODEL_PATH}. Pulando o download.")
-
-# --- Fim: Bloco de Download Automático do Modelo ---
+# --- Fim do Bloco de Download ---
 
 
 from image_datasets.dataset import image_resize
@@ -50,16 +42,15 @@ def generate(image: Image.Image, edit_prompt: str):
     from src.flux.xflux_pipeline import XFluxSampler
     
     global sampler
-    if sampler == None:
-        # A inicialização do sampler agora ocorrerá após a confirmação de que o modelo foi baixado.
+    if sampler is None:
+        # CORREÇÃO: Inicializa o sampler usando os argumentos do arquivo .yaml
+        print("Inicializando o XFluxSampler com a configuração...")
         sampler = XFluxSampler(
-            device = device,
-            ip_loaded=False,
-            spatial_condition=False,
-            clip_image_processor=None,
-            image_encoder=None,
-            improj=None,
-            share_position_embedding = True,
+            device=device,
+            ip_loaded=args.use_ip,
+            spatial_condition=args.use_spatial_condition,
+            share_position_embedding=args.share_position_embedding,
+            use_share_weight_referencenet=args.use_share_weight_referencenet
         )
 
     img = image_resize(image, 544)
@@ -68,6 +59,9 @@ def generate(image: Image.Image, edit_prompt: str):
     img = torch.from_numpy((np.array(img) / 127.5) - 1)
     img = img.permute(2, 0, 1).unsqueeze(0).to(device, dtype=dtype)
 
+    # CORREÇÃO: Passa a imagem de origem se qualquer modo de condicionamento estiver ativo
+    use_image_conditioning = args.use_spatial_condition or args.use_share_weight_referencenet
+    
     result = sampler(
         prompt=edit_prompt,
         width=args.sample_width,
@@ -75,9 +69,9 @@ def generate(image: Image.Image, edit_prompt: str):
         num_steps=args.sample_steps,
         image_prompt=None,
         true_gs=args.cfg_scale,
-        seed=args.seed,
+        seed=args.seed if args.seed != -1 else np.random.randint(0, 2**32 - 1),
         ip_scale=args.ip_scale if args.use_ip else 1.0,
-        source_image=img if args.use_spatial_condition else None,
+        source_image=img if use_image_conditioning else None,
     )
     return result
 
@@ -201,7 +195,6 @@ def create_app():
             </div>
             """
         )
-        # gr.Markdown(header, elem_id="header")
         with gr.Row(equal_height=False):
             with gr.Column(variant="panel", elem_classes="inputPanel"):
                 original_image = gr.Image(