fixing bugs

app.py

@@ -1,195 +1,129 @@
-# app.py
 import gradio as gr
 import torch
-import jax
-import jax.numpy as jnp
 import numpy as np
-from PIL import Image
+import jax
 import pickle
-import logging
+from PIL import Image
 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, 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()]
+from super_resolve import process
+from diffusers import StableDiffusionXLPipeline
+from transformers import DPTFeatureExtractor, DPTForDepthEstimation
+
+# ========== Configuração do Thera ==========
+REPO_ID_EDSR = "prs-eth/thera-edsr-pro"
+REPO_ID_RDN = "prs-eth/thera-rdn-pro"
+
+
+# Carregar modelos Thera
+def load_thera_model(repo_id):
+    model_path = hf_hub_download(repo_id=repo_id, filename="model.pkl")
+    with open(model_path, 'rb') as fh:
+        check = pickle.load(fh)
+        params, backbone, size = check['model'], check['backbone'], check['size']
+        model = build_thera(3, backbone, size)
+    return model, params
+
+
+model_edsr, params_edsr = load_thera_model(REPO_ID_EDSR)
+model_rdn, params_rdn = load_thera_model(REPO_ID_RDN)
+
+# ========== Configuração do SDXL + Depth ==========
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
+# Carregar modelos de geração
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch_dtype
+).to(device)
+
+pipe.load_lora_weights(
+    "KappaNeuro/bas-relief",
+    weight_name="BAS-RELIEF.safetensors",
+    peft_backend="peft"
 )
-logger = logging.getLogger(__name__)
 
-# Configurações
-JAX_DEVICE = jax.devices("cpu")[0]
-TORCH_DEVICE = "cpu"
-
-
-def load_thera_model(repo_id: str, filename: str):
-    """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 segura
-try:
-    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 crítica na inicialização: {str(e)}")
-    raise
-
-
-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 de erros robusto"""
-    try:
-        # 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 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
-        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}"
-            )
-
-        # Pré-processamento da imagem
-        source = jnp.array(image).astype(jnp.float32) / 255.0
-        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))
-
-        # 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}")
-
-        # 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,
-            guidance_scale=7.5
-        )
-        bas_relief = result.images[0]
-        logger.info(f"Bas-Relief gerado: {bas_relief.size}")
-
-        # 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 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"Falha no processamento: {str(e)}")
-
-
-# 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():
-        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],
-        api_name="processar"
+feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
+depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
+
+
+# ========== Funções Principais ==========
+def super_resolution(image, scale_factor, model_type):
+    model = model_edsr if model_type == "EDSR" else model_rdn
+    params = params_edsr if model_type == "EDSR" else params_rdn
+
+    source = np.asarray(image) / 255.
+    target_shape = (
+        round(source.shape[0] * scale_factor),
+        round(source.shape[1] * scale_factor),
+    )
+
+    output = process(source, model, params, target_shape, do_ensemble=True)
+    return Image.fromarray(np.asarray(output))
+
+
+def generate_bas_relief(prompt):
+    full_prompt = f"BAS-RELIEF {prompt}"
+    image = pipe(
+        prompt=full_prompt,
+        num_inference_steps=25,
+        guidance_scale=7.5,
+        height=512,
+        width=512
+    ).images[0]
+
+    inputs = feature_extractor(image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = depth_model(**inputs)
+        depth_map = outputs.predicted_depth
+
+    depth_map = torch.nn.functional.interpolate(
+        depth_map.unsqueeze(1),
+        size=image.size[::-1],
+        mode="bicubic"
+    ).squeeze().cpu().numpy()
+
+    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
+    depth_map = (depth_map * 255).astype(np.uint8)
+
+    return image, Image.fromarray(depth_map)
+
+
+# ========== Interface Gradio ==========
+with gr.Blocks(title="TheraSR + Bas-Relief Generator") as app:
+    gr.Markdown("# 🔥 TheraSR + Bas-Relief Generator")
+    gr.Markdown("Combine aliasing-free super-resolution with artistic bas-relief generation")
+
+    with gr.Tabs():
+        with gr.TabItem("🖼 Super-Resolution"):
+            with gr.Row():
+                sr_input = gr.Image(label="Input Image", type="pil")
+                sr_output = gr.Image(label="Super-Resolution Result")
+            sr_scale = gr.Slider(1.0, 6.0, value=2.0, label="Scale Factor")
+            sr_model = gr.Radio(["EDSR", "RDN"], value="EDSR", label="Model Type")
+            sr_btn = gr.Button("Enhance Resolution")
+
+        with gr.TabItem("🎨 Generate Bas-Relief"):
+            with gr.Row():
+                text_input = gr.Textbox(label="Art Prompt", placeholder="Roman soldier marble relief...")
+            with gr.Row():
+                gen_output = gr.Image(label="Generated Art")
+                depth_output = gr.Image(label="Depth Map")
+            gen_btn = gr.Button("Generate Artwork")
+
+    # Event Handlers
+    sr_btn.click(
+        super_resolution,
+        inputs=[sr_input, sr_scale, sr_model],
+        outputs=sr_output
+    )
+
+    gen_btn.click(
+        generate_bas_relief,
+        inputs=text_input,
+        outputs=[gen_output, depth_output]
     )
 
-if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+# Configuração do Hugging Face
+app.launch(debug=False, share=True)

requirements.txt

@@ -1,39 +1,11 @@
--f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-
-ConfigArgParse==1.7
-Pillow==10.0.0
-chex==0.1.7
-diffusers
-einops==0.6.1
-flax==0.6.10
-flaxmodels==0.1.3
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.0.1
+torchvision==0.15.2
+jax[cuda11_pip]==0.4.13
+flax==0.7.4
+diffusers==0.24.0
+transformers==4.35.2
+peft==0.6.2
+gradio==4.12.0
 huggingface_hub
-jax==0.4.11
-jaxlib==0.4.11+cuda11.cudnn86
-jaxtyping==0.2.20
-ml-dtypes==0.1.0
-numpy==1.24.1
-nvidia-cublas-cu11==11.11.3.6
-nvidia-cuda-cupti-cu11==11.8.87
-nvidia-cuda-nvcc-cu11==11.8.89
-nvidia-cuda-runtime-cu11==11.8.89
-nvidia-cudnn-cu11==8.9.2.26
-nvidia-cufft-cu11==10.9.0.58
-nvidia-cusolver-cu11==11.4.1.48
-nvidia-cusparse-cu11==11.7.5.86
-opt-einsum==3.3.0
-optax==0.2.0
-orbax-checkpoint==0.2.4
-peft
-Pillow==10.0.0
-scipy==1.10.1
-timm==0.9.6
-torch
-torchvision
-tqdm==4.65.0
-transformers==4.46.3
-wandb
-
-gradio==4.44.1
-gradio_imageslider==0.0.20
-spaces
+pillow==10.0.0

utils.py

@@ -1,69 +1,36 @@
-# utils.py
+from functools import partial
+
 import jax
-import jax.numpy as jnp
 import numpy as np
-from functools import partial
 
 
-def repeat_vmap(fun, in_axes=None):
-    if in_axes is None:
-        in_axes = [0]
+def repeat_vmap(fun, in_axes=[0]):
     for axes in in_axes:
         fun = jax.vmap(fun, in_axes=axes)
     return fun
 
 
-def make_grid(target_shape: tuple[int, int]):
-    """Gera grid de coordenadas com validação rigorosa"""
-    h, w = target_shape
-
-    # 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)
-
-    # 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 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):
-    """Interpolação segura com verificações em tempo real"""
-    try:
-        # Converter e garantir formato 4D
-        coords = jnp.asarray(coords)
-        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"Formato inválido: {original_shape} → {coords.shape}. "
-                f"Esperado (B, H, W, 2)"
-            )
-
-        # Transformação de coordenadas
-        coords = coords.transpose((0, 3, 1, 2))
-        coords = coords.at[:, 0].set(
-            coords[:, 0] * (grid.shape[-3] - 1) + (grid.shape[-3] - 1) / 2
-        )
-        coords = coords.at[:, 1].set(
-            coords[:, 1] * (grid.shape[-2] - 1) + (grid.shape[-2] - 1) / 2
-        )
-
-        # Interpolação vetorizada
-        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 na interpolação: {str(e)}") from e
+    """
+    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)