fixing bugs

app.py

@@ -1,106 +1,106 @@
-import logging
-import pickle
-import warnings
-
+# app.py
 import gradio as gr
+import torch
 import jax
 import jax.numpy as jnp
 import numpy as np
-import torch
 from PIL import Image
-from diffusers import StableDiffusionXLImg2ImgPipeline
+import pickle
+import logging
 from huggingface_hub import hf_hub_download
+from diffusers import StableDiffusionXLImg2ImgPipeline
 from transformers import DPTImageProcessor, DPTForDepthEstimation
-
 from model import build_thera
-from utils import make_grid
+from utils import make_grid, interpolate_grid
 
 # Configuração de logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.FileHandler("processing.log"),
-        logging.StreamHandler()
-    ]
+    handlers=[logging.FileHandler("processing.log"), logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
 # Configurações
-warnings.filterwarnings("ignore")
 JAX_DEVICE = jax.devices("cpu")[0]
 TORCH_DEVICE = "cpu"
 
 
-def load_thera_model(repo_id, filename):
+def load_thera_model(repo_id: str, filename: str):
+    """Carrega modelo com verificação de segurança"""
     try:
         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']
-        return build_thera(3, backbone, size), variables
+            checkpoint = pickle.load(fh)
+        return build_thera(3, checkpoint['backbone'], checkpoint['size']), checkpoint['model']
     except Exception as e:
-        logger.error(f"Erro ao carregar Thera: {str(e)}")
+        logger.error(f"Erro ao carregar modelo: {str(e)}")
         raise
 
 
-logger.info("Carregando modelos...")
-model_edsr, variables_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")
-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")
-feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
-depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(TORCH_DEVICE)
+# Inicialização dos modelos
+try:
+    logger.info("Carregando modelos...")
+    model_edsr, params_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")
+    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")
+    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)}")
+    raise
 
 
-def adjust_size(size):
-    return max(8, (size // 8) * 8)
+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 full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
+def full_pipeline(image: Image.Image, prompt: str, scale_factor: float = 2.0):
+    """Pipeline completo com tratamento robusto"""
     try:
-        progress(0.1, desc="Iniciando...")
+        # Pré-processamento
         image = image.convert("RGB")
-        source = np.array(image) / 255.0
+        orig_w, orig_h = image.size
 
-        # Ajuste de dimensões
-        target_shape = (
-            adjust_size(int(image.height * scale_factor)),
-            adjust_size(int(image.width * scale_factor))
-        )
-        logger.info(f"Transformação: {image.size} → {target_shape}")
+        # 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}")
 
-        # Gerar grid
-        coords = make_grid(target_shape)
-        logger.debug(f"Coords shape: {coords.shape}")
+        # 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})")
 
         # Super-resolução
-        progress(0.3, desc="Processando super-resolução...")
-        source_jax = jax.device_put(source[np.newaxis, ...], JAX_DEVICE)
+        source = jnp.array(image).astype(jnp.float32) / 255.0
+        source = source[jnp.newaxis, ...]  # Adicionar batch
+
         t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)
+        upscaled = model_edsr.apply(params_edsr, source, t, (new_h, new_w))
 
-        upscaled = model_edsr.apply(
-            variables_edsr,
-            source_jax,
-            t,
-            target_shape
-        )
-        upscaled_pil = Image.fromarray((np.array(upscaled[0]) * 255).astype(np.uint8))
+        # Pós-processamento
+        upscaled_img = Image.fromarray((np.array(upscaled[0]) * 255).astype(np.uint8))
 
         # Bas-Relief
-        progress(0.6, desc="Gerando relevo...")
-        bas_relief = pipe(
-            prompt=f"BAS-RELIEF {prompt}, ultra detailed engraving, 16K resolution",
-            image=upscaled_pil,
+        result = pipe(
+            prompt=f"BAS-RELIEF {prompt}, ultra detailed, 8K resolution",
+            image=upscaled_img,
             strength=0.7,
-            num_inference_steps=25
-        ).images[0]
+            num_inference_steps=30
+        )
+        bas_relief = result.images[0]
 
-        # Depth Map
-        progress(0.8, desc="Calculando profundidade...")
+        # Mapa de profundidade
         inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
         with torch.no_grad():
             depth = depth_model(**inputs).predicted_depth
@@ -111,30 +111,49 @@ def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
             mode="bicubic"
         ).squeeze().cpu().numpy()
 
-        depth_normalized = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-        depth_pil = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+        # Normalizaçã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))
 
-        return upscaled_pil, bas_relief, depth_pil
+        return upscaled_img, bas_relief, depth_img
 
     except Exception as e:
-        logger.error(f"ERRO: {str(e)}", exc_info=True)
-        raise gr.Error(f"Erro no processamento: {str(e)}")
+        logger.error(f"ERRO NO PIPELINE: {str(e)}", exc_info=True)
+        raise gr.Error(f"Processamento falhou: {str(e)}")
 
 
 # Interface
-with gr.Blocks(title="SuperRes + BasRelief") as app:
-    gr.Markdown("## 🖼️ Super Resolução + 🗿 Bas-Relief + 🗺️ Mapa de Profundidade")
+with gr.Blocks(title="SuperRes+BasRelief", theme=gr.themes.Default()) as app:
+    gr.Markdown("# 🖼️ Super Resolução + 🗿 Bas-Relief + 🗺️ Mapa de Profundidade")
+
     with gr.Row():
         with gr.Column():
-            img_input = gr.Image(type="pil", label="Entrada")
-            prompt = gr.Textbox("Escultura detalhada em mármore, alto relevo", label="Descrição")
-            scale = gr.Slider(1.0, 4.0, value=2.0, label="Escala")
-            btn = gr.Button("Processar ▶️")
+            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():
-            img_upscaled = gr.Image(label="Super Resolução")
-            img_basrelief = gr.Image(label="Bas-Relief")
-            img_depth = gr.Image(label="Profundidade")
-    btn.click(full_pipeline, [img_input, prompt, scale], [img_upscaled, img_basrelief, img_depth])
+            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(
+        full_pipeline,
+        inputs=[img_input, prompt, scale],
+        outputs=[upscaled_output, basrelief_output, depth_output]
+    )
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch(server_name="0.0.0.0", server_port=7860)

utils.py

@@ -1,40 +1,40 @@
-from functools import partial
+# utils.py
 import jax
 import jax.numpy as jnp
 import numpy as np
-
-
-def repeat_vmap(fun, in_axes=[0]):
-    for axes in in_axes:
-        fun = jax.vmap(fun, in_axes=axes)
-    return fun
+from functools import partial
 
 
 def make_grid(patch_size: int | tuple[int, int]):
-    """Gera grid de coordenadas com segurança numérica"""
-    # Garantir tamanho mínimo de 8x8
+    """Gera grid de coordenadas com validação robusta"""
     if isinstance(patch_size, int):
-        h = w = max(8, patch_size)
+        h = w = max(16, patch_size)  # Novo mínimo seguro
     else:
-        h, w = (max(8, ps) for ps in patch_size)
+        h, w = (max(16, ps) for ps in patch_size)  # 16x16 mínimo
 
-    # Espaçamento preciso entre pontos
-    y_space = np.linspace(-0.5 + 1 / (2 * h), 0.5 - 1 / (2 * h), h)
-    x_space = np.linspace(-0.5 + 1 / (2 * w), 0.5 - 1 / (2 * w), w)
+    # Cálculo preciso das coordenadas
+    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)
 
-    # Criar grid com dimensões (1, H, W, 2)
-    grid = np.stack(np.meshgrid(y_space, x_space, indexing='ij'), axis=-1)
+    # Grid com dimensões (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 segura com verificação de dimensões"""
+    """Interpolação com tratamento completo de dimensões"""
     try:
-        # Converter para JAX array e validar formato
+        # Converter e garantir 4D
         coords = jnp.asarray(coords)
-        if coords.ndim != 4 or coords.shape[-1] != 2:
+        if coords.ndim == 1:  # Caso de erro reportado
+            coords = coords.reshape(1, 1, 1, -1)
+        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}. Esperado (B, H, W, 2)"
+                f"Dimensões inválidas: {coords.shape}. Formato esperado: (B, H, W, 2)"
             )
 
         # Transformação de coordenadas
@@ -47,12 +47,10 @@ def interpolate_grid(coords, grid, order=0):
         )
 
         # Interpolação vetorizada
-        map_fn = jax.vmap(jax.vmap(
-            partial(jax.scipy.ndimage.map_coordinates, order=order, mode='nearest'),
-            in_axes=(2, None),
-            out_axes=2
-        ))
-        return map_fn(grid, coords)
+        map_coordinates = partial(jax.scipy.ndimage.map_coordinates,
+                                  order=order,
+                                  mode='nearest')
+        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