ds1david commited on
Commit
98889c8
·
1 Parent(s): d9330a6

First commit

Browse files
Files changed (2) hide show
  1. app.py +96 -0
  2. requirements.txt +9 -0
app.py ADDED
@@ -0,0 +1,96 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import numpy as np
4
+ from diffusers import StableDiffusionXLImg2ImgPipeline
5
+ from transformers import DPTFeatureExtractor, DPTForDepthEstimation
6
+ from PIL import Image, ImageEnhance, ImageOps
7
+
8
+ device = "cpu" # or "cuda" if you have a GPU
9
+ torch_dtype = torch.float32
10
+
11
+ print("Loading SDXL Img2Img model...")
12
+ pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
13
+ "stabilityai/stable-diffusion-xl-base-1.0",
14
+ torch_dtype=torch_dtype
15
+ ).to(device)
16
+
17
+ print("Loading bas-relief LoRA weights with PEFT...")
18
+ pipe.load_lora_weights(
19
+ "KappaNeuro/bas-relief",
20
+ weight_name="BAS-RELIEF.safetensors",
21
+ peft_backend="peft"
22
+ )
23
+
24
+ print("Loading DPT Depth Model...")
25
+ feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
26
+ depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
27
+
28
+
29
+ def enhance_depth_map(depth_arr: np.ndarray) -> Image.Image:
30
+ d_min, d_max = depth_arr.min(), depth_arr.max()
31
+ depth_stretched = (depth_arr - d_min) / (d_max - d_min + 1e-8)
32
+ depth_stretched = (depth_stretched * 255).astype(np.uint8)
33
+
34
+ depth_pil = Image.fromarray(depth_stretched)
35
+ depth_pil = ImageOps.autocontrast(depth_pil)
36
+
37
+ enhancer = ImageEnhance.Sharpness(depth_pil)
38
+ depth_pil = enhancer.enhance(2.0)
39
+
40
+ return depth_pil
41
+
42
+
43
+ def generate_bas_relief_and_depth(input_image: Image.Image):
44
+ # Redimensionar a imagem para o tamanho esperado
45
+ input_image = input_image.resize((512, 512))
46
+
47
+ # Prompt fixo para ativar o LoRA
48
+ prompt = "BAS-RELIEF"
49
+
50
+ print("Gerando imagem no estilo baixo-relevo...")
51
+ result = pipe(
52
+ prompt=prompt,
53
+ image=input_image,
54
+ strength=0.7, # Controla a intensidade da transformação
55
+ num_inference_steps=15,
56
+ guidance_scale=7.5
57
+ )
58
+ generated_image = result.images[0]
59
+
60
+ print("Calculando mapa de profundidade...")
61
+ inputs = feature_extractor(generated_image, return_tensors="pt").to(device)
62
+ with torch.no_grad():
63
+ outputs = depth_model(**inputs)
64
+ predicted_depth = outputs.predicted_depth
65
+
66
+ prediction = torch.nn.functional.interpolate(
67
+ predicted_depth.unsqueeze(1),
68
+ size=generated_image.size[::-1],
69
+ mode="bicubic",
70
+ align_corners=False
71
+ ).squeeze()
72
+
73
+ depth_map_pil = enhance_depth_map(prediction.cpu().numpy())
74
+
75
+ return generated_image, depth_map_pil
76
+
77
+
78
+ title = "Conversor para Baixo-relevo (SDXL + LoRA) com Mapa de Profundidade"
79
+ description = (
80
+ "Carrega stable-diffusion-xl-base-1.0 no CPU, aplica LoRA de 'KappaNeuro/bas-relief' "
81
+ "para transformar imagens em baixo-relevo e calcula o mapa de profundidade correspondente."
82
+ )
83
+
84
+ iface = gr.Interface(
85
+ fn=generate_bas_relief_and_depth,
86
+ inputs=gr.Image(label="Imagem de Entrada", type="pil"),
87
+ outputs=[
88
+ gr.Image(label="Imagem em Baixo-relevo"),
89
+ gr.Image(label="Mapa de Profundidade")
90
+ ],
91
+ title=title,
92
+ description=description
93
+ )
94
+
95
+ if __name__ == "__main__":
96
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ open3d
2
+ peft
3
+ accelerate
4
+ diffusers>=0.20.0
5
+ transformers>=4.30.0
6
+ torch
7
+ gradio
8
+ Pillow
9
+ safetensors