app.py

import logging
import pickle
import warnings

import gradio as gr
import jax
import jax.numpy as jnp
import numpy as np
import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline
from huggingface_hub import hf_hub_download
from transformers import DPTImageProcessor, DPTForDepthEstimation

from model import build_thera
from utils import make_grid

# 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
warnings.filterwarnings("ignore")
JAX_DEVICE = jax.devices("cpu")[0]
TORCH_DEVICE = "cpu"


def load_thera_model(repo_id, filename):
    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
    except Exception as e:
        logger.error(f"Erro ao carregar Thera: {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)


def adjust_size(size):
    return max(8, (size // 8) * 8)


def full_pipeline(image, prompt, scale_factor=2.0, progress=gr.Progress()):
    try:
        progress(0.1, desc="Iniciando...")
        image = image.convert("RGB")
        source = np.array(image) / 255.0

        # 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}")

        # Gerar grid
        coords = make_grid(target_shape)
        logger.debug(f"Coords shape: {coords.shape}")

        # Super-resolução
        progress(0.3, desc="Processando super-resolução...")
        source_jax = jax.device_put(source[np.newaxis, ...], JAX_DEVICE)
        t = jnp.array([1.0 / (scale_factor ** 2)], dtype=jnp.float32)

        upscaled = model_edsr.apply(
            variables_edsr,
            source_jax,
            t,
            target_shape
        )
        upscaled_pil = 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,
            strength=0.7,
            num_inference_steps=25
        ).images[0]

        # Depth Map
        progress(0.8, desc="Calculando profundidade...")
        inputs = feature_extractor(bas_relief, return_tensors="pt").to(TORCH_DEVICE)
        with torch.no_grad():
            depth = depth_model(**inputs).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 no processamento: {str(e)}")


# Interface
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="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 ▶️")
        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])

if __name__ == "__main__":
    app.launch()