Trying fix variants

app.py

@@ -1,33 +1,31 @@
 import gradio as gr
 import torch
 import numpy as np
-from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers import StableDiffusionXLPipeline
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 from PIL import Image, ImageEnhance, ImageOps
 
-# Configuração de dispositivo
 device = "cpu"  # or "cuda" if you have a GPU
 torch_dtype = torch.float32
 
-print("Carregando modelo SDXL Img2Img...")
-pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+print("Loading SDXL Base model...")
+pipe = StableDiffusionXLPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch_dtype
 ).to(device)
 
-print("Carregando pesos LoRA weights with PEFT...")
+print("Loading bas-relief LoRA weights with PEFT...")
 pipe.load_lora_weights(
-    "KappaNeuro/bas-relief",
+    "KappaNeuro/bas-relief",      # The HF repo with BAS-RELIEF.safetensors
     weight_name="BAS-RELIEF.safetensors",
-    peft_backend="peft"
+    peft_backend="peft"          # This is crucial
 )
 
-print("Carregando modelo de profundidade...")
+print("Loading DPT Depth Model...")
 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:
+def enhance_depth_map(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)
@@ -40,16 +38,16 @@ def processar_profundidade(depth_arr: np.ndarray) -> Image.Image:
 
     return depth_pil
 
-
-def processar_imagem(imagem: Image.Image):
-    # Pré-processamento
+def generate_bas_relief_and_depth(imagem):
+    # Use the token "BAS-RELIEF" so the LoRA triggers
+    full_prompt = f"BAS-RELIEF {prompt}"
     print("Generating image with LoRA style...")
     result = pipe(
-        prompt="BAS-RELIEF",
+        prompt=full_prompt,
         image=imagem,
-        num_inference_steps=15,  # reduce if too slow
+        num_inference_steps=15,   # reduce if too slow
         guidance_scale=7.5,
-        height=512,  # reduce if you still get timeouts
+        height=512,               # reduce if you still get timeouts
         width=512
     )
     image = result.images[0]
@@ -67,16 +65,13 @@ def processar_imagem(imagem: Image.Image):
         align_corners=False
     ).squeeze()
 
-    depth_map_pil = processar_profundidade(prediction.cpu().numpy())
+    depth_map_pil = enhance_depth_map(prediction.cpu().numpy())
 
     return image, depth_map_pil
 
-    # return resultado.images[0], processar_profundidade(depth_map)
-
-
 # Interface Gradio
 interface = gr.Interface(
-    fn=processar_imagem,
+    fn=generate_bas_relief_and_depth,
     inputs=gr.Image(type="pil"),
     outputs=[gr.Image(label="Resultado"), gr.Image(label="Profundidade")],
     title="Conversor para Baixo-relevo",