New logic

app.py

@@ -6,11 +6,24 @@ import numpy as np
 from PIL import Image
 import pickle
 import warnings
+import logging
+from datetime import datetime
 from huggingface_hub import hf_hub_download
 from diffusers import StableDiffusionXLImg2ImgPipeline
 from transformers import DPTImageProcessor, DPTForDepthEstimation
 from model import build_thera
 
+# Configuração de logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("processing.log"),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
 # Configurações e supressão de avisos
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=UserWarning)
@@ -22,56 +35,86 @@ TORCH_DEVICE = "cpu"
 
 # 1. Carregar modelos do Thera ----------------------------------------------------------------
 def load_thera_model(repo_id, filename):
-    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
-    with open(model_path, 'rb') as fh:
-        check = pickle.load(fh)
-        # Carregar estrutura completa de variáveis
-        variables = check['model']  # Deve conter {'params': ...}
-        backbone, size = check['backbone'], check['size']
-    model = build_thera(3, backbone, size)
-    return model, variables
+    try:
+        logger.info(f"Iniciando carregamento do modelo Thera de {repo_id}")
+        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        with open(model_path, 'rb') as fh:
+            check = pickle.load(fh)
+            variables = check['model']
+            backbone, size = check['backbone'], check['size']
+        model = build_thera(3, backbone, size)
+        logger.info("Modelo Thera carregado com sucesso")
+        return model, variables
+    except Exception as e:
+        logger.error(f"Falha ao carregar modelo Thera: {str(e)}")
+        raise
 
 
-print("Carregando Thera EDSR...")
+logger.info("Carregando Thera EDSR...")
 model_edsr, variables_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")
 
 # 2. Carregar SDXL + LoRA ---------------------------------------------------------------------
-print("Carregando SDXL + LoRA...")
-pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float32
-).to(TORCH_DEVICE)
-pipe.load_lora_weights("KappaNeuro/bas-relief", weight_name="BAS-RELIEF.safetensors")
+try:
+    logger.info("Iniciando carregamento do SDXL + LoRA...")
+    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        torch_dtype=torch.float32
+    ).to(TORCH_DEVICE)
+    pipe.load_lora_weights("KappaNeuro/bas-relief", weight_name="BAS-RELIEF.safetensors")
+    logger.info("SDXL + LoRA carregado com sucesso")
+except Exception as e:
+    logger.error(f"Falha ao carregar SDXL: {str(e)}")
+    raise
 
 # 3. Carregar modelo de profundidade ----------------------------------------------------------
-print("Carregando DPT Depth...")
-feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
-depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(TORCH_DEVICE)
+try:
+    logger.info("Iniciando carregamento do DPT Depth...")
+    feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
+    depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(TORCH_DEVICE)
+    logger.info("Modelo DPT carregado com sucesso")
+except Exception as e:
+    logger.error(f"Falha ao carregar DPT: {str(e)}")
+    raise
 
 
 # Pipeline principal --------------------------------------------------------------------------
-def full_pipeline(image, prompt, scale_factor=2.0):
+def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
     try:
+        progress(0, desc="Iniciando processamento...")
+
         # 1. Super Resolução com Thera
+        progress(0.1, desc="Convertendo imagem para RGB...")
         image = image.convert("RGB")
+
+        progress(0.2, desc="Preparando entrada para super-resolução...")
         source = np.array(image) / 255.0
+        original_size = image.size
         target_shape = (int(image.height * scale_factor), int(image.width * scale_factor))
+        logger.info(f"Super-resolução: {original_size} → {target_shape} (scale: {scale_factor}x)")
 
+        progress(0.3, desc="Processando com Thera...")
         source_jax = jax.device_put(source, JAX_DEVICE)
         t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)
 
-        # Chamada corrigida com estrutura de variáveis correta
+        start_time = datetime.now()
         upscaled = model_edsr.apply(
-            variables_edsr,  # Estrutura completa {'params': ...}
+            variables_edsr,
             source_jax,
             t,
             target_shape
         )
+        logger.info(f"Super-resolução concluída em {datetime.now() - start_time}")
 
+        progress(0.5, desc="Convertendo resultado...")
         upscaled_pil = Image.fromarray((np.array(upscaled) * 255).astype(np.uint8))
+        logger.info(f"Tamanho após super-resolução: {upscaled_pil.size}")
 
         # 2. Gerar Bas-Relief
+        progress(0.6, desc="Gerando Bas-Relief...")
         full_prompt = f"BAS-RELIEF {prompt}, insanely detailed and complex engraving relief, ultra-high definition, rich in detail, 16K resolution"
+        logger.info(f"Prompt final: {full_prompt}")
+
+        start_time = datetime.now()
         bas_relief = pipe(
             prompt=full_prompt,
             image=upscaled_pil,
@@ -79,9 +122,13 @@ def full_pipeline(image, prompt, scale_factor=2.0):
             num_inference_steps=25,
             guidance_scale=7.5
         ).images[0]
+        logger.info(f"Bas-Relief gerado em {datetime.now() - start_time}")
 
         # 3. Calcular Depth Map
+        progress(0.8, desc="Calculando mapa de profundidade...")
+        start_time = datetime.now()
         inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
+
         with torch.no_grad():
             outputs = depth_model(**inputs)
             depth = outputs.predicted_depth
@@ -92,12 +139,18 @@ def full_pipeline(image, prompt, scale_factor=2.0):
             mode="bicubic"
         ).squeeze().cpu().numpy()
 
-        depth_normalized = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
+        progress(0.9, desc="Processando mapa de profundidade...")
+        depth_min = depth_map.min()
+        depth_max = depth_map.max()
+        depth_normalized = (depth_map - depth_min) / (depth_max - depth_min + 1e-8)
         depth_pil = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+        logger.info(f"Profundidade calculada em {datetime.now() - start_time} | Range: {depth_min:.2f}-{depth_max:.2f}")
 
+        progress(1.0, desc="Finalizado!")
         return upscaled_pil, bas_relief, depth_pil
 
     except Exception as e:
+        logger.error(f"Erro no processamento: {str(e)}", exc_info=True)
         raise gr.Error(f"Erro no processamento: {str(e)}")
 
 
@@ -127,4 +180,5 @@ with gr.Blocks(title="Super Res + Bas-Relief") as app:
     )
 
 if __name__ == "__main__":
+    logger.info("Iniciando aplicação Gradio")
     app.launch(share=False)