fixing bugs

app.py

@@ -27,132 +27,168 @@ TORCH_DEVICE = "cpu"
 
 
 def load_thera_model(repo_id: str, filename: str):
-    """Carrega modelo com verificação de segurança"""
+    """Carrega modelo com múltiplas verificações"""
     try:
         model_path = hf_hub_download(repo_id=repo_id, filename=filename)
         with open(model_path, 'rb') as fh:
             checkpoint = pickle.load(fh)
+
+        # Verificar estrutura do checkpoint
+        required_keys = {'model', 'backbone', 'size'}
+        if not required_keys.issubset(checkpoint.keys()):
+            missing = required_keys - checkpoint.keys()
+            raise ValueError(f"Checkpoint corrompido. Chaves faltando: {missing}")
+
         return build_thera(3, checkpoint['backbone'], checkpoint['size']), checkpoint['model']
     except Exception as e:
         logger.error(f"Erro ao carregar modelo: {str(e)}")
         raise
 
 
-# Inicialização dos modelos
+# Inicialização segura
 try:
-    logger.info("Carregando modelos...")
+    logger.info("Inicializando modelos...")
     model_edsr, params_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")
+
+    # Pipeline SDXL
     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")
+
+    # Modelo de profundidade
     feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
     depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(TORCH_DEVICE)
+
 except Exception as e:
-    logger.error(f"Falha na inicialização: {str(e)}")
+    logger.error(f"Falha crítica na inicialização: {str(e)}")
     raise
 
 
-def adjust_size(original: int, scale: float) -> int:
-    """Ajuste de tamanho com limites seguros"""
-    scaled = int(original * scale)
-    adjusted = (scaled // 8) * 8  # Divisível por 8
-    return max(32, adjusted)  # Mínimo absoluto
+def safe_resize(original: tuple[int, int], scale: float) -> tuple[int, int]:
+    """Calcula tamanho garantindo estabilidade numérica"""
+    h, w = original
+    new_h = int(h * scale)
+    new_w = int(w * scale)
+
+    # Ajustar para múltiplo de 8
+    new_h = max(32, new_h - new_h % 8)
+    new_w = max(32, new_w - new_w % 8)
+
+    return (new_h, new_w)
 
 
 def full_pipeline(image: Image.Image, prompt: str, scale_factor: float = 2.0):
-    """Pipeline completo com tratamento robusto"""
+    """Pipeline completo com tratamento de erros robusto"""
     try:
-        # Pré-processamento
+        # Verificação inicial
+        if not image:
+            raise ValueError("Nenhuma imagem fornecida")
+
+        # Conversão segura para RGB
         image = image.convert("RGB")
         orig_w, orig_h = image.size
+        logger.info(f"Processando imagem: {orig_w}x{orig_h}")
 
-        # Cálculo do tamanho alvo
-        new_h = adjust_size(orig_h, scale_factor)
-        new_w = adjust_size(orig_w, scale_factor)
-        logger.info(f"Redimensionando: {orig_h}x{orig_w} → {new_h}x{new_w}")
+        # Cálculo do novo tamanho
+        new_h, new_w = safe_resize((orig_h, orig_w), scale_factor)
+        logger.info(f"Novo tamanho calculado: {new_h}x{new_w}")
 
         # Gerar grid de coordenadas
-        coords = make_grid((new_h, new_w))
-        logger.debug(f"Dimensões do grid: {coords.shape}")
-
-        # Verificação crítica
-        if coords.shape[1:3] != (new_h, new_w):
-            raise ValueError(f"Grid incorreto: {coords.shape[1:3]} vs ({new_h}, {new_w})")
+        grid = make_grid((new_h, new_w))
+        logger.debug(f"Grid gerado: {grid.shape}")
+
+        # Verificação crítica do grid
+        if grid.shape[1:3] != (new_h, new_w):
+            raise RuntimeError(
+                f"Incompatibilidade de dimensões: "
+                f"Grid {grid.shape[1:3]} vs Alvo {new_h}x{new_w}"
+            )
 
-        # Super-resolução
+        # Pré-processamento da imagem
         source = jnp.array(image).astype(jnp.float32) / 255.0
-        source = source[jnp.newaxis, ...]  # Adicionar batch
+        source = source[jnp.newaxis, ...]  # Adicionar dimensão de batch
 
+        # Parâmetro de escala
         t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)
+
+        # Processamento Thera
         upscaled = model_edsr.apply(params_edsr, source, t, (new_h, new_w))
 
-        # Pós-processamento
+        # Conversão para PIL
         upscaled_img = Image.fromarray((np.array(upscaled[0]) * 255).astype(np.uint8))
+        logger.info(f"Imagem super-resolvida: {upscaled_img.size}")
 
-        # Bas-Relief
+        # Geração do Bas-Relief
         result = pipe(
             prompt=f"BAS-RELIEF {prompt}, ultra detailed, 8K resolution",
             image=upscaled_img,
             strength=0.7,
-            num_inference_steps=30
+            num_inference_steps=30,
+            guidance_scale=7.5
         )
         bas_relief = result.images[0]
+        logger.info(f"Bas-Relief gerado: {bas_relief.size}")
 
-        # Mapa de profundidade
+        # Cálculo da profundidade
         inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
         with torch.no_grad():
             depth = depth_model(**inputs).predicted_depth
 
+        # Redimensionamento
         depth_map = torch.nn.functional.interpolate(
             depth.unsqueeze(1),
             size=bas_relief.size[::-1],
             mode="bicubic"
         ).squeeze().cpu().numpy()
 
-        # Normalização
+        # Normalização e conversão
         depth_min = depth_map.min()
         depth_max = depth_map.max()
         depth_normalized = (depth_map - depth_min) / (depth_max - depth_min + 1e-8)
         depth_img = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+        logger.info("Mapa de profundidade calculado")
 
         return upscaled_img, bas_relief, depth_img
 
     except Exception as e:
         logger.error(f"ERRO NO PIPELINE: {str(e)}", exc_info=True)
-        raise gr.Error(f"Processamento falhou: {str(e)}")
+        raise gr.Error(f"Falha no processamento: {str(e)}")
 
 
-# Interface
-with gr.Blocks(title="SuperRes+BasRelief", theme=gr.themes.Default()) as app:
+# Interface Gradio
+with gr.Blocks(title="SuperRes+BasRelief Pro", theme=gr.themes.Soft()) as app:
     gr.Markdown("# 🖼️ Super Resolução + 🗿 Bas-Relief + 🗺️ Mapa de Profundidade")
 
     with gr.Row():
-        with gr.Column():
-            img_input = gr.Image(label="Imagem de Entrada", type="pil")
-            prompt = gr.Textbox(
-                label="Descrição do Relevo",
-                value="Ainsanely detailed and complex engraving relief, ultra-high definition",
-                placeholder="Descreva o estilo desejado..."
-            )
-            scale = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
-            btn = gr.Button("Processar Imagem", variant="primary")
-
-        with gr.Column():
-            gr.Markdown("## Resultados")
-            with gr.Tabs():
-                with gr.TabItem("Super Resolução"):
-                    upscaled_output = gr.Image(label="Resultado Super Resolução")
-                with gr.TabItem("Bas-Relief"):
-                    basrelief_output = gr.Image(label="Relevo Gerado")
-                with gr.TabItem("Profundidade"):
-                    depth_output = gr.Image(label="Mapa de Profundidade")
-
-    btn.click(
+        input_col = gr.Column()
+        output_col = gr.Column()
+
+    with input_col:
+        img_input = gr.Image(label="Carregar Imagem", type="pil", height=300)
+        prompt = gr.Textbox(
+            label="Descrição do Relevo",
+            value="A insanely detailed and complex engraving relief, ultra-high definition",
+            placeholder="Descreva o estilo desejado..."
+        )
+        scale = gr.Slider(1.0, 4.0, value=2.0, step=0.1, label="Fator de Escala")
+        process_btn = gr.Button("Iniciar Processamento", variant="primary")
+
+    with output_col:
+        with gr.Tabs():
+            with gr.TabItem("Super Resolução"):
+                upscaled_output = gr.Image(label="Resultado", show_label=False)
+            with gr.TabItem("Bas-Relief"):
+                basrelief_output = gr.Image(label="Relevo", show_label=False)
+            with gr.TabItem("Profundidade"):
+                depth_output = gr.Image(label="Mapa 3D", show_label=False)
+
+    process_btn.click(
         full_pipeline,
         inputs=[img_input, prompt, scale],
-        outputs=[upscaled_output, basrelief_output, depth_output]
+        outputs=[upscaled_output, basrelief_output, depth_output],
+        api_name="processar"
     )
 
 if __name__ == "__main__":

utils.py

@@ -13,36 +13,41 @@ def repeat_vmap(fun, in_axes=None):
     return fun
 
 
-def make_grid(patch_size: int | tuple[int, int]):
-    """Gera grid de coordenadas com validação robusta"""
-    if isinstance(patch_size, int):
-        h = w = max(16, patch_size)  # Novo mínimo seguro
-    else:
-        h, w = (max(16, ps) for ps in patch_size)  # 16x16 mínimo
+def make_grid(target_shape: tuple[int, int]):
+    """Gera grid de coordenadas com validação rigorosa"""
+    h, w = target_shape
 
-    # Cálculo preciso das coordenadas
+    # Garantir tamanho mínimo e divisibilidade
+    h = max(32, h)
+    w = max(32, w)
+    h = h if h % 8 == 0 else h + (8 - h % 8)
+    w = w if w % 8 == 0 else w + (8 - w % 8)
+
+    # Espaçamento preciso
     y_coords = np.linspace(-0.5 + 1 / (2 * h), 0.5 - 1 / (2 * h), h)
     x_coords = np.linspace(-0.5 + 1 / (2 * w), 0.5 - 1 / (2 * w), w)
 
-    # Grid com dimensões (1, H, W, 2)
+    # Criar grid 4D (1, H, W, 2)
     grid = np.stack(np.meshgrid(y_coords, x_coords, indexing='ij'), axis=-1)
     return grid[np.newaxis, ...]
 
 
 def interpolate_grid(coords, grid, order=0):
-    """Interpolação com tratamento completo de dimensões"""
+    """Interpolação segura com verificações em tempo real"""
     try:
-        # Converter e garantir 4D
+        # Converter e garantir formato 4D
         coords = jnp.asarray(coords)
-        if coords.ndim == 1:  # Caso de erro reportado
-            coords = coords.reshape(1, 1, 1, -1)
+        original_shape = coords.shape
+
+        # Adicionar dimensões faltantes
         while coords.ndim < 4:
             coords = coords[jnp.newaxis, ...]
 
         # Validação final
         if coords.shape[-1] != 2 or coords.ndim != 4:
             raise ValueError(
-                f"Dimensões inválidas: {coords.shape}. Formato esperado: (B, H, W, 2)"
+                f"Formato inválido: {original_shape} → {coords.shape}. "
+                f"Esperado (B, H, W, 2)"
             )
 
         # Transformação de coordenadas
@@ -61,4 +66,4 @@ def interpolate_grid(coords, grid, order=0):
         return jax.vmap(jax.vmap(map_coordinates, in_axes=(2, None), out_axes=2))(grid, coords)
 
     except Exception as e:
-        raise RuntimeError(f"Erro de interpolação: {str(e)}") from e
+        raise RuntimeError(f"Erro na interpolação: {str(e)}") from e