app.py

import gradio as gr
import torch
import numpy as np
from diffusers import StableDiffusionXLImg2ImgPipeline
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from PIL import Image, ImageEnhance, ImageOps

# Configuração de dispositivo
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float32 if device == "cpu" else torch.float16

print("Carregando modelo SDXL Img2Img...")
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch_dtype,
    use_safetensors=True
).to(device)

print("Carregando pesos LoRA...")
pipe.load_lora_weights(
    "KappaNeuro/bas-relief",
    weight_name="BAS-RELIEF.safetensors",
    adapter_name="bas_relief"
)

print("Carregando modelo de profundidade...")
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)


def processar_profundidade(depth_arr: np.ndarray) -> Image.Image:
    depth_normalized = (depth_arr - depth_arr.min()) / (depth_arr.max() - depth_arr.min() + 1e-8)
    depth_img = Image.fromarray((depth_normalized * 255).astype(np.uint8))
    return ImageOps.autocontrast(depth_img)


def processar_imagem(imagem: Image.Image):
    # Pré-processamento
    imagem = imagem.convert("RGB").resize((512, 512))

    # Gerar baixo-relevo
    with torch.inference_mode():
        resultado = pipe(
            prompt="BAS-RELIEF",
            image=imagem,
            strength=0.7,
            num_inference_steps=20,
            guidance_scale=7.5
        )

    # Calcular profundidade
    inputs = feature_extractor(resultado.images[0], return_tensors="pt").to(device)
    with torch.no_grad():
        depth = depth_model(**inputs).predicted_depth

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

    return resultado.images[0], processar_profundidade(depth_map)


# Interface Gradio
interface = gr.Interface(
    fn=processar_imagem,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(label="Resultado"), gr.Image(label="Profundidade")],
    title="Conversor para Baixo-relevo",
    description="Transforme imagens em baixo-relevo com mapa de profundidade"
)

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