Trying fix variants

app.py

@@ -5,98 +5,71 @@ from diffusers import StableDiffusionXLImg2ImgPipeline
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 from PIL import Image, ImageEnhance, ImageOps
 
-# Configuração de dispositivo e tipos de dados
+# Configuração de dispositivo
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
+torch_dtype = torch.float32 if device == "cpu" else torch.float16
 
 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...")
+print("Carregando pesos LoRA...")
 pipe.load_lora_weights(
     "KappaNeuro/bas-relief",
     weight_name="BAS-RELIEF.safetensors",
-    adapter_name="bas_relief",
-    peft_backend="peft"
+    adapter_name="bas_relief"
 )
 
-print("Carregando modelo de profundidade DPT...")
+print("Carregando modelo de profundidade...")
 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)
+def processar_profundidade(depth_arr: np.ndarray) -> Image.Image:
+    depth_normalized = (depth_arr - depth_arr.min()) / (depth_arr.max() - depth_arr.min() + 1e-8)
+    depth_img = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+    return ImageOps.autocontrast(depth_img)
 
-    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
+def processar_imagem(imagem: Image.Image):
+    # Pré-processamento
     imagem = imagem.convert("RGB").resize((512, 512))
 
-    # Geração da imagem em baixo-relevo
-    with torch.autocast(device, dtype=torch_dtype):
+    # Gerar baixo-relevo
+    with torch.inference_mode():
         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)
+            num_inference_steps=20,
+            guidance_scale=7.5
         )
 
-    imagem_gerada = resultado.images[0]
-
-    # Cálculo do mapa de profundidade
-    inputs = feature_extractor(imagem_gerada, return_tensors="pt").to(device)
+    # Calcular profundidade
+    inputs = feature_extractor(resultado.images[0], return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = depth_model(**inputs)
-        predicted_depth = outputs.predicted_depth
+        depth = depth_model(**inputs).predicted_depth
 
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=imagem_gerada.size[::-1],
-        mode="bicubic",
-        align_corners=False
-    ).squeeze()
+    depth_map = torch.nn.functional.interpolate(
+        depth.unsqueeze(1),
+        size=imagem.size[::-1],
+        mode="bicubic"
+    ).squeeze().cpu().numpy()
 
-    mapa_profundidade = melhorar_mapa_profundidade(prediction.cpu().numpy())
-
-    return imagem_gerada, mapa_profundidade
+    return resultado.images[0], processar_profundidade(depth_map)
 
 
 # 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"
+    fn=processar_imagem,
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Image(label="Resultado"), gr.Image(label="Profundidade")],
+    title="Conversor para Baixo-relevo",
+    description="Transforme imagens em baixo-relevo com mapa de profundidade"
 )
 
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0" if torch.cuda.is_available() else None)
+    interface.launch()