Trying fix variants

app.py

@@ -5,28 +5,32 @@ from diffusers import StableDiffusionXLImg2ImgPipeline
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 from PIL import Image, ImageEnhance, ImageOps
 
-device = "cpu"  # or "cuda" if you have a GPU
-torch_dtype = torch.float32
+# 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("Loading SDXL Img2Img model...")
+print("Carregando modelo SDXL Img2Img...")
 pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch_dtype
+    torch_dtype=torch_dtype,
+    variant="fp32",
+    use_safetensors=True
 ).to(device)
 
-print("Loading bas-relief LoRA weights with PEFT...")
+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("Loading DPT Depth Model...")
+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 enhance_depth_map(depth_arr: np.ndarray) -> Image.Image:
+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)
@@ -40,57 +44,59 @@ def enhance_depth_map(depth_arr: np.ndarray) -> Image.Image:
     return depth_pil
 
 
-def generate_bas_relief_and_depth(input_image: Image.Image):
-    # Redimensionar a imagem para o tamanho esperado
-    input_image = input_image.resize((512, 512))
+def gerar_baixo_relevo_e_profundidade(imagem: Image.Image):
+    # Pré-processamento da imagem
+    imagem = imagem.convert("RGB").resize((512, 512))
 
-    # Prompt fixo para ativar o LoRA
-    prompt = "BAS-RELIEF"
+    # 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)
+        )
 
-    print("Gerando imagem no estilo baixo-relevo...")
-    result = pipe(
-        prompt=prompt,
-        image=input_image,
-        strength=0.7,  # Controla a intensidade da transformação
-        num_inference_steps=15,
-        guidance_scale=7.5
-    )
-    generated_image = result.images[0]
+    imagem_gerada = resultado.images[0]
 
-    print("Calculando mapa de profundidade...")
-    inputs = feature_extractor(generated_image, return_tensors="pt").to(device)
+    # 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=generated_image.size[::-1],
+        size=imagem_gerada.size[::-1],
         mode="bicubic",
         align_corners=False
     ).squeeze()
 
-    depth_map_pil = enhance_depth_map(prediction.cpu().numpy())
+    mapa_profundidade = melhorar_mapa_profundidade(prediction.cpu().numpy())
 
-    return generated_image, depth_map_pil
+    return imagem_gerada, mapa_profundidade
 
 
-title = "Conversor para Baixo-relevo (SDXL + LoRA) com Mapa de Profundidade"
-description = (
-    "Carrega stable-diffusion-xl-base-1.0 no CPU, aplica LoRA de 'KappaNeuro/bas-relief' "
-    "para transformar imagens em baixo-relevo e calcula o mapa de profundidade correspondente."
+# 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."
 )
 
-iface = gr.Interface(
-    fn=generate_bas_relief_and_depth,
+interface = gr.Interface(
+    fn=gerar_baixo_relevo_e_profundidade,
     inputs=gr.Image(label="Imagem de Entrada", type="pil"),
     outputs=[
-        gr.Image(label="Imagem em Baixo-relevo"),
+        gr.Image(label="Baixo-relevo Gerado"),
         gr.Image(label="Mapa de Profundidade")
     ],
-    title=title,
-    description=description
+    title=titulo,
+    description=descricao,
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch(server_name="0.0.0.0" if torch.cuda.is_available() else None)