File size: 3,269 Bytes
98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a 98889c8 e87f39a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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 e tipos de dados
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):
# Pré-processamento da imagem
imagem = imagem.convert("RGB").resize((512, 512))
# Geração da imagem em baixo-relevo
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]
# Cálculo do mapa de profundidade
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
# Interface Gradio
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) |