|
|
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 |
|
|
|
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
torch_dtype = torch.float16 if device == "cuda" else torch.float32 |
|
|
|
|
|
print("Carregando modelo SDXL Img2Img...") |
|
|
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( |
|
|
"stabilityai/stable-diffusion-xl-base-1.0", |
|
|
torch_dtype=torch_dtype, |
|
|
variant="fp32", |
|
|
use_safetensors=True |
|
|
).to(device) |
|
|
|
|
|
print("Carregando pesos LoRA para baixo-relevo...") |
|
|
pipe.load_lora_weights( |
|
|
"KappaNeuro/bas-relief", |
|
|
weight_name="BAS-RELIEF.safetensors", |
|
|
adapter_name="bas_relief", |
|
|
peft_backend="peft" |
|
|
) |
|
|
|
|
|
print("Carregando modelo de profundidade DPT...") |
|
|
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large") |
|
|
depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device) |
|
|
|
|
|
|
|
|
def melhorar_mapa_profundidade(depth_arr: np.ndarray) -> Image.Image: |
|
|
d_min, d_max = depth_arr.min(), depth_arr.max() |
|
|
depth_stretched = (depth_arr - d_min) / (d_max - d_min + 1e-8) |
|
|
depth_stretched = (depth_stretched * 255).astype(np.uint8) |
|
|
|
|
|
depth_pil = Image.fromarray(depth_stretched) |
|
|
depth_pil = ImageOps.autocontrast(depth_pil) |
|
|
|
|
|
enhancer = ImageEnhance.Sharpness(depth_pil) |
|
|
depth_pil = enhancer.enhance(2.0) |
|
|
|
|
|
return depth_pil |
|
|
|
|
|
|
|
|
def gerar_baixo_relevo_e_profundidade(imagem: Image.Image): |
|
|
|
|
|
imagem = imagem.convert("RGB").resize((512, 512)) |
|
|
|
|
|
|
|
|
with torch.autocast(device, dtype=torch_dtype): |
|
|
resultado = pipe( |
|
|
prompt="BAS-RELIEF", |
|
|
image=imagem, |
|
|
strength=0.7, |
|
|
num_inference_steps=15, |
|
|
guidance_scale=7.5, |
|
|
generator=torch.Generator(device=device).manual_seed(0) |
|
|
) |
|
|
|
|
|
imagem_gerada = resultado.images[0] |
|
|
|
|
|
|
|
|
inputs = feature_extractor(imagem_gerada, return_tensors="pt").to(device) |
|
|
with torch.no_grad(): |
|
|
outputs = depth_model(**inputs) |
|
|
predicted_depth = outputs.predicted_depth |
|
|
|
|
|
prediction = torch.nn.functional.interpolate( |
|
|
predicted_depth.unsqueeze(1), |
|
|
size=imagem_gerada.size[::-1], |
|
|
mode="bicubic", |
|
|
align_corners=False |
|
|
).squeeze() |
|
|
|
|
|
mapa_profundidade = melhorar_mapa_profundidade(prediction.cpu().numpy()) |
|
|
|
|
|
return imagem_gerada, mapa_profundidade |
|
|
|
|
|
|
|
|
|
|
|
titulo = "Conversor para Baixo-relevo com Mapa de Profundidade" |
|
|
descricao = ( |
|
|
"Carrega uma imagem para transformar em estilo baixo-relevo usando SDXL + LoRA " |
|
|
"e gera o mapa de profundidade correspondente." |
|
|
) |
|
|
|
|
|
interface = gr.Interface( |
|
|
fn=gerar_baixo_relevo_e_profundidade, |
|
|
inputs=gr.Image(label="Imagem de Entrada", type="pil"), |
|
|
outputs=[ |
|
|
gr.Image(label="Baixo-relevo Gerado"), |
|
|
gr.Image(label="Mapa de Profundidade") |
|
|
], |
|
|
title=titulo, |
|
|
description=descricao, |
|
|
allow_flagging="never" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
interface.launch(server_name="0.0.0.0" if torch.cuda.is_available() else None) |
