app.py

# app.py
import logging
import gradio as gr
import torch
import numpy as np
import jax
import pickle
from PIL import Image
from huggingface_hub import hf_hub_download
from model import build_thera
from super_resolve import process
from diffusers import StableDiffusionXLImg2ImgPipeline
from transformers import DPTFeatureExtractor, DPTForDepthEstimation


# ================== CONFIGURAÇÃO DE LOGGING ==================
class CustomLogger:
    def __init__(self, name):
        self.logger = logging.getLogger(name)
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.setLevel(logging.INFO)


    def divider(self, text=None, length=60):
        if text:
            # Cálculo seguro do número de '='
            available_space = length - len(text) - 12  # 10 '=' + 2 espaços
            if available_space < 1:
                available_space = 1  # Garante pelo menos 1 '='
            msg = f"\n{'=' * 10} {text.upper()} {'=' * available_space}"
        else:
            msg = "\n" + "=" * length
        self.logger.info(msg)


    def etapa(self, text):
        self.logger.info(f"▶ {text}")

    def success(self, text):
        self.logger.info(f"✓ {text}")

    def error(self, text):
        self.logger.error(f"✗ {text}")

    def warning(self, text):
        self.logger.warning(f"⚠ {text}")


logger = CustomLogger(__name__)

# ================== CONFIGURAÇÃO DE HARDWARE ==================
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
logger.divider("Configuração Inicial")
logger.success(f"Dispositivo detectado: {device.upper()}")
logger.success(f"Precisão numérica: {str(torch_dtype).replace('torch.', '')}")


# ================== CARREGAMENTO DE MODELOS ==================
def carregar_modelo_thera(repo_id):
    try:
        logger.divider(f"Carregando Modelo: {repo_id}")
        model_path = hf_hub_download(repo_id=repo_id, filename="model.pkl")
        with open(model_path, 'rb') as f:
            check = pickle.load(f)
            model = build_thera(3, check['backbone'], check['size'])
            params = check['model']
        logger.success(f"Modelo {repo_id} carregado")
        return model, params
    except Exception as e:
        logger.error(f"Falha ao carregar {repo_id}: {str(e)}")
        return None, None


# Carregar modelos Thera
try:
    modelo_edsr, params_edsr = carregar_modelo_thera("prs-eth/thera-edsr-pro")
    modelo_rdn, params_rdn = carregar_modelo_thera("prs-eth/thera-rdn-pro")
except Exception as e:
    logger.error("Falha crítica no carregamento dos modelos Thera")
    raise

# ================== PIPELINE DE ARTE ==================
pipe = None
modelo_profundidade = None

try:
    logger.divider("Configurando Pipeline de Arte")
    pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch_dtype,
        variant="fp16",
        use_safetensors=True
    ).to(device)

    pipe.load_lora_weights(
        "KappaNeuro/bas-relief",
        weight_name="BAS-RELIEF.safetensors"
    )
    logger.success("Pipeline SDXL configurado")

    logger.etapa("Configurando Modelo de Profundidade")
    processador_profundidade = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
    modelo_profundidade = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
    logger.success("Modelo de profundidade pronto")

except Exception as e:
    logger.error(f"Erro na configuração da GPU: {str(e)}")
    pipe = None
    modelo_profundidade = None


# ================== FLUXO DE PROCESSAMENTO ==================
def pipeline_completo(imagem, fator_escala, modelo_escolhido, prompt_estilo):
    try:
        logger.divider("Novo Processamento")

        # Converter entrada
        if not isinstance(imagem, Image.Image):
            imagem = Image.fromarray(imagem)

        # ========= SUPER-RESOLUÇÃO =========
        logger.etapa("Processando Super-Resolução")
        modelo = modelo_edsr if modelo_escolhido == "EDSR" else modelo_rdn
        params = params_edsr if modelo_escolhido == "EDSR" else params_rdn

        sr_array = process(
            np.array(imagem) / 255.,
            modelo,
            params,
            (round(imagem.size[1] * fator_escala),
             round(imagem.size[0] * fator_escala)),
            True
        )

        sr_pil = Image.fromarray(np.array(sr_array)).convert("RGB")
        logger.success(f"Super-Resolução: {sr_pil.size[0]}x{sr_pil.size[1]}")

        # ========= ESTILO BAIXO-RELEVO =========
        arte_pil = None
        if pipe and modelo_profundidade:
            try:
                logger.etapa("Aplicando Estilo Artístico")
                resultado = pipe(
                    prompt=f"BAS-RELIEF {prompt_estilo}, intricate marble carving, 8k ultra HD",
                    image=sr_pil,
                    strength=0.65,
                    num_inference_steps=30,
                    guidance_scale=7.5
                )
                arte_pil = resultado.images[0]
                logger.success(f"Arte gerada: {arte_pil.size[0]}x{arte_pil.size[1]}")
            except Exception as e:
                logger.error(f"Falha no estilo: {str(e)}")

        # ========= MAPA DE PROFUNDIDADE =========
        mapa_pil = None
        if arte_pil and modelo_profundidade:
            try:
                logger.etapa("Calculando Profundidade")
                inputs = processador_profundidade(images=arte_pil, return_tensors="pt").to(device)
                with torch.no_grad():
                    outputs = modelo_profundidade(**inputs)
                    depth = outputs.predicted_depth

                depth = torch.nn.functional.interpolate(
                    depth.unsqueeze(1).float(),
                    size=arte_pil.size[::-1],
                    mode="bicubic"
                ).squeeze().cpu().numpy()

                depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
                mapa_pil = Image.fromarray((depth * 255).astype(np.uint8))
                logger.success("Mapa de profundidade calculado")
            except Exception as e:
                logger.error(f"Falha na profundidade: {str(e)}")

        return sr_pil, arte_pil or sr_pil, mapa_pil or sr_pil

    except Exception as e:
        logger.error(f"Erro no pipeline: {str(e)}")
        return None, None, None


# ================== INTERFACE GRADIO ==================
with gr.Blocks(title="TheraSR Art Suite", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🎨 TheraSR - Super Resolução & Arte Generativa")

    with gr.Row():
        with gr.Column():
            entrada_imagem = gr.Image(label="Imagem de Entrada", type="pil")
            seletor_modelo = gr.Radio(
                ["EDSR", "RDN"],
                value="EDSR",
                label="Modelo de Super-Resolução"
            )
            controle_escala = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
            entrada_prompt = gr.Textbox(
                label="Prompt de Estilo",
                value="insanely detailed ancient greek marble浮雕, 8k cinematic lighting"
            )
            botao_processar = gr.Button("Gerar", variant="primary")

        with gr.Column():
            saida_sr = gr.Image(label="Super-Resolução", show_label=True)
            saida_arte = gr.Image(label="Arte em Relevo", show_label=True)
            saida_profundidade = gr.Image(label="Mapa de Profundidade", show_label=True)

    botao_processar.click(
        pipeline_completo,
        inputs=[entrada_imagem, controle_escala, seletor_modelo, entrada_prompt],
        outputs=[saida_sr, saida_arte, saida_profundidade]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)