fixing bugs

app.py

@@ -7,7 +7,6 @@ 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
@@ -36,17 +35,16 @@ TORCH_DEVICE = "cpu"
 # 1. Carregar modelos do Thera ----------------------------------------------------------------
 def load_thera_model(repo_id, filename):
     try:
-        logger.info(f"Iniciando carregamento do modelo Thera de {repo_id}")
+        logger.info(f"Carregando 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)}")
+        logger.error(f"Erro ao carregar modelo: {str(e)}")
         raise
 
 
@@ -55,80 +53,78 @@ model_edsr, variables_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.p
 
 # 2. Carregar SDXL + LoRA ---------------------------------------------------------------------
 try:
-    logger.info("Iniciando carregamento do SDXL + LoRA...")
+    logger.info("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")
-    logger.info("SDXL + LoRA carregado com sucesso")
 except Exception as e:
-    logger.error(f"Falha ao carregar SDXL: {str(e)}")
+    logger.error(f"Erro ao carregar SDXL: {str(e)}")
     raise
 
 # 3. Carregar modelo de profundidade ----------------------------------------------------------
 try:
-    logger.info("Iniciando carregamento do DPT Depth...")
+    logger.info("Carregando 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)}")
+    logger.error(f"Erro ao carregar DPT: {str(e)}")
     raise
 
 
-# Pipeline principal --------------------------------------------------------------------------
+def adjust_size(size):
+    """Garante que o tamanho seja divisível por 8"""
+    return (size // 8) * 8
+
+
 def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
     try:
-        progress(0, desc="Iniciando processamento...")
+        progress(0.1, desc="Pré-processamento...")
 
-        # 1. Super Resolução com Thera
-        progress(0.1, desc="Convertendo imagem para RGB...")
+        # Converter e verificar imagem
         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...")
+        # Adicionar dimensão de batch se necessário
+        if source.ndim == 3:
+            source = source[np.newaxis, ...]
+
+        # Ajustar tamanho alvo
+        target_shape = (
+            adjust_size(int(image.height * scale_factor)),
+            adjust_size(int(image.width * scale_factor))
+        )
+
+        progress(0.3, desc="Super-resolução...")
         source_jax = jax.device_put(source, JAX_DEVICE)
         t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)
 
-        start_time = datetime.now()
+        # Processar com Thera
         upscaled = model_edsr.apply(
             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...")
+        # Remover dimensão de batch se necessário
+        if upscaled.ndim == 4:
+            upscaled = upscaled[0]
+
         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()
+        full_prompt = f"BAS-RELIEF {prompt}, ultra detailed engraving, 16K resolution"
         bas_relief = pipe(
             prompt=full_prompt,
             image=upscaled_pil,
             strength=0.7,
-            num_inference_steps=25,
-            guidance_scale=7.5
+            num_inference_steps=25
         ).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()
+        progress(0.8, desc="Calculando profundidade...")
         inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
-
         with torch.no_grad():
             outputs = depth_model(**inputs)
             depth = outputs.predicted_depth
@@ -139,39 +135,34 @@ def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
             mode="bicubic"
         ).squeeze().cpu().numpy()
 
-        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_normalized = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
         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)}")
+        logger.error(f"Erro: {str(e)}", exc_info=True)
+        raise gr.Error(f"Erro: {str(e)}")
 
 
 # Interface Gradio ----------------------------------------------------------------------------
-with gr.Blocks(title="Super Res + Bas-Relief") as app:
-    gr.Markdown("## 🔍 Super Resolução + 🗿 Bas-Relief + 🗺️ Profundidade")
+with gr.Blocks(title="SuperRes + BasRelief") 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="Imagem de Entrada")
             prompt = gr.Textbox(
-                label="Descrição do Relevo",
-                value="insanely detailed and complex engraving relief, ultra-high definition, rich in detail, and 16K resolution."
+                label="Descrição",
+                value="insanely detailed and complex engraving relief, ultra-high definition"
             )
             scale = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
             btn = gr.Button("Processar")
 
         with gr.Column():
-            img_upscaled = gr.Image(label="Imagem Super Resolvida")
-            img_basrelief = gr.Image(label="Resultado Bas-Relief")
-            img_depth = gr.Image(label="Mapa de Profundidade")
+            img_upscaled = gr.Image(label="Super Resolvida")
+            img_basrelief = gr.Image(label="Bas-Relief")
+            img_depth = gr.Image(label="Profundidade")
 
     btn.click(
         full_pipeline,
@@ -180,5 +171,4 @@ with gr.Blocks(title="Super Res + Bas-Relief") as app:
     )
 
 if __name__ == "__main__":
-    logger.info("Iniciando aplicação Gradio")
-    app.launch(share=False)
+    app.launch()  # Sem compartilhamento público

utils.py

@@ -1,174 +1,36 @@
-import gradio as gr
-import torch
+from functools import partial
+
 import jax
-import jax.numpy as jnp
 import numpy as np
-from PIL import Image
-import pickle
-import warnings
-import logging
-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)
-
-# Configurar dispositivos
-JAX_DEVICE = jax.devices("cpu")[0]
-TORCH_DEVICE = "cpu"
-
-
-# 1. Carregar modelos do Thera ----------------------------------------------------------------
-def load_thera_model(repo_id, filename):
-    try:
-        logger.info(f"Carregando 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)
-        return model, variables
-    except Exception as e:
-        logger.error(f"Erro ao carregar modelo: {str(e)}")
-        raise
-
-
-logger.info("Carregando Thera EDSR...")
-model_edsr, variables_edsr = load_thera_model("prs-eth/thera-edsr-pro", "model.pkl")
-
-# 2. Carregar SDXL + LoRA ---------------------------------------------------------------------
-try:
-    logger.info("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")
-except Exception as e:
-    logger.error(f"Erro ao carregar SDXL: {str(e)}")
-    raise
-
-# 3. Carregar modelo de profundidade ----------------------------------------------------------
-try:
-    logger.info("Carregando DPT Depth...")
-    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"Erro ao carregar DPT: {str(e)}")
-    raise
-
-
-def adjust_size(size):
-    """Garante que o tamanho seja divisível por 8"""
-    return (size // 8) * 8
-
-
-def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
-    try:
-        progress(0.1, desc="Pré-processamento...")
-
-        # Converter e verificar imagem
-        image = image.convert("RGB")
-        source = np.array(image) / 255.0
-
-        # Adicionar dimensão de batch se necessário
-        if source.ndim == 3:
-            source = source[np.newaxis, ...]
-
-        # Ajustar tamanho alvo
-        target_shape = (
-            adjust_size(int(image.height * scale_factor)),
-            adjust_size(int(image.width * scale_factor))
-        )
-
-        progress(0.3, desc="Super-resolução...")
-        source_jax = jax.device_put(source, JAX_DEVICE)
-        t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)
-
-        # Processar com Thera
-        upscaled = model_edsr.apply(
-            variables_edsr,
-            source_jax,
-            t,
-            target_shape
-        )
-
-        # Remover dimensão de batch se necessário
-        if upscaled.ndim == 4:
-            upscaled = upscaled[0]
-
-        upscaled_pil = Image.fromarray((np.array(upscaled) * 255).astype(np.uint8))
-
-        progress(0.6, desc="Gerando Bas-Relief...")
-        full_prompt = f"BAS-RELIEF {prompt}, ultra detailed engraving, 16K resolution"
-        bas_relief = pipe(
-            prompt=full_prompt,
-            image=upscaled_pil,
-            strength=0.7,
-            num_inference_steps=25
-        ).images[0]
-
-        progress(0.8, desc="Calculando profundidade...")
-        inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
-        with torch.no_grad():
-            outputs = depth_model(**inputs)
-            depth = outputs.predicted_depth
-
-        depth_map = torch.nn.functional.interpolate(
-            depth.unsqueeze(1),
-            size=bas_relief.size[::-1],
-            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))
-
-        return upscaled_pil, bas_relief, depth_pil
-
-    except Exception as e:
-        logger.error(f"Erro: {str(e)}", exc_info=True)
-        raise gr.Error(f"Erro: {str(e)}")
-
-
-# Interface Gradio ----------------------------------------------------------------------------
-with gr.Blocks(title="SuperRes + BasRelief") 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="Imagem de Entrada")
-            prompt = gr.Textbox(
-                label="Descrição",
-                value="insanely detailed and complex engraving relief, ultra-high definition"
-            )
-            scale = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
-            btn = gr.Button("Processar")
-
-        with gr.Column():
-            img_upscaled = gr.Image(label="Super Resolvida")
-            img_basrelief = gr.Image(label="Bas-Relief")
-            img_depth = gr.Image(label="Profundidade")
 
-    btn.click(
-        full_pipeline,
-        inputs=[img_input, prompt, scale],
-        outputs=[img_upscaled, img_basrelief, img_depth]
-    )
 
-if __name__ == "__main__":
-    app.launch()  # Sem compartilhamento público
+def repeat_vmap(fun, in_axes=[0]):
+    for axes in in_axes:
+        fun = jax.vmap(fun, in_axes=axes)
+    return fun
+
+
+def make_grid(patch_size: int | tuple[int, int]):
+    if isinstance(patch_size, int):
+        patch_size = (patch_size, patch_size)
+    offset_h, offset_w = 1 / (2 * np.array(patch_size))
+    space_h = np.linspace(-0.5 + offset_h, 0.5 - offset_h, patch_size[0])
+    space_w = np.linspace(-0.5 + offset_w, 0.5 - offset_w, patch_size[1])
+    return np.stack(np.meshgrid(space_h, space_w, indexing='ij'), axis=-1)  # [h, w]
+
+
+def interpolate_grid(coords, grid, order=0):
+    """
+    args:
+        coords: Tensor of shape (B, H, W, 2) with coordinates in [-0.5, 0.5]
+        grid: Tensor of shape (B, H', W', C)
+    returns:
+        Tensor of shape (B, H, W, C) with interpolated values
+    """
+    # convert [-0.5, 0.5] -> [0, size], where pixel centers are expected at
+    # [-0.5 + 1 / (2*size), ..., 0.5 - 1 / (2*size)]
+    coords = coords.transpose((0, 3, 1, 2))
+    coords = coords.at[:, 0].set(coords[:, 0] * grid.shape[-3] + (grid.shape[-3] - 1) / 2)
+    coords = coords.at[:, 1].set(coords[:, 1] * grid.shape[-2] + (grid.shape[-2] - 1) / 2)
+    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)