fixing bugs

app.py

@@ -40,23 +40,19 @@ class CustomLogger:
     def error(self, text):
         self.logger.error(f"✗ {text}")
 
-    def warning(self, text):
-        self.logger.warning(f"⚠ {text}")
-
 
 logger = CustomLogger(__name__)
 
-# ================== CONFIGURAÇÃO DE HARDWARE ==================
+# ================== CONFIGURAÇÃO FORÇADA ==================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float32  # Forçar precisão única para compatibilidade
-logger.divider("Inicialização do Sistema")
-logger.success(f"Dispositivo detectado: {device.upper()}")
-logger.success(f"Modo de precisão: float32")
+torch_dtype = torch.float32  # Forçar float32 universalmente
+logger.divider("Configuração Forçada")
+logger.success(f"Dispositivo: {device.upper()}")
+logger.success(f"Precisão: {str(torch_dtype).replace('torch.', '')}")
 
 
 # ================== CARREGAMENTO DE MODELOS ==================
 def carregar_modelo_thera(repo_id):
-    """Carrega modelos Thera com tratamento de erros robusto"""
     try:
         logger.divider(f"Carregando {repo_id}")
         model_path = hf_hub_download(repo_id=repo_id, filename="model.pkl")
@@ -67,30 +63,33 @@ def carregar_modelo_thera(repo_id):
         logger.success(f"{repo_id} carregado")
         return model, params
     except Exception as e:
-        logger.error(f"Falha ao carregar {repo_id}: {str(e)}")
+        logger.error(f"Falha no carregamento: {str(e)}")
         return None, None
 
 
-# Carregar modelos principais
-modelo_edsr, params_edsr = carregar_modelo_thera("prs-eth/thera-edsr-pro")
-modelo_rdn, params_rdn = carregar_modelo_thera("prs-eth/thera-rdn-pro")
+try:
+    modelo_edsr, params_edsr = carregar_modelo_thera("prs-eth/thera-edsr-pro")
+    modelo_rdn, params_rdn = carregar_modelo_thera("prs-eth/thera-rdn-pro")
+except Exception as e:
+    logger.error("Falha crítica nos modelos Thera")
+    raise
 
-# ================== MODELOS DE ARTE (CARREGAMENTO CONDICIONAL) ==================
+# ================== PIPELINE ARTÍSTICO ==================
 pipe = None
 modelo_profundidade = None
-processador_profundidade = None
 
 try:
-    logger.divider("Inicializando Componentes Artísticos")
+    logger.divider("Configurando Componentes Artísticos")
 
-    # Pipeline de estilo
+    # Pipeline principal
     pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0",
         torch_dtype=torch_dtype,
-        use_safetensors=True
+        use_safetensors=True,
+        variant=None  # Forçar carregamento sem fp16
     ).to(device)
 
-    # Adapter LoRA
+    # LoRA
     pipe.load_lora_weights(
         "KappaNeuro/bas-relief",
         weight_name="BAS-RELIEF.safetensors"
@@ -98,61 +97,60 @@ try:
 
     # Modelo de profundidade
     processador_profundidade = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
-    modelo_profundidade = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
+    modelo_profundidade = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device).float()
 
-    logger.success("Componentes artísticos prontos")
+    logger.success("Componentes artísticos em float32")
 except Exception as e:
-    logger.warning(f"Recursos artísticos desativados: {str(e)}")
+    logger.warning(f"Recursos artísticos limitados: {str(e)}")
     pipe = None
-    modelo_profundidade = None
 
 
-# ================== PIPELINE PRINCIPAL ==================
-def processar_imagem(imagem, escala, modelo, prompt):
-    """Fluxo completo de processamento com fallbacks"""
+# ================== PROCESSAMENTO PRINCIPAL ==================
+def processar_imagem_completa(imagem, escala, modelo, prompt):
     try:
-        logger.divider("Novo Processamento")
+        logger.divider("Iniciando Processamento")
 
-        # Converter entrada para PIL
+        # Converter entrada
         if not isinstance(imagem, Image.Image):
             imagem = Image.fromarray(imagem)
 
         # ========= 1. SUPER-RESOLUÇÃO =========
-        logger.etapa("Super-Resolução Thera")
+        logger.etapa("Processando Super-Resolução")
         modelo_sr = modelo_edsr if modelo == "EDSR" else modelo_rdn
         params_sr = params_edsr if modelo == "EDSR" else params_rdn
 
         sr_jax = process(
-            np.array(imagem) / 255.0,
+            np.array(imagem) / 255.,
             modelo_sr,
             params_sr,
-            (int(imagem.height * escala),
-             int(imagem.width * escala)),
+            (round(imagem.height * escala),
+             round(imagem.width * escala)),
             True
         )
 
         sr_pil = Image.fromarray(np.array(sr_jax)).convert("RGB")
-        logger.success(f"Resolução: {sr_pil.size[0]}x{sr_pil.size[1]}")
+        logger.success(f"SR: {sr_pil.size[0]}x{sr_pil.size[1]}")
 
         # ========= 2. ESTILO BAIXO-RELEVO =========
-        arte_pil = sr_pil  # Fallback padrão
+        arte_pil = sr_pil  # Fallback
         if pipe:
             try:
                 logger.etapa("Aplicando Estilo")
                 arte_pil = pipe(
-                    prompt=f"BAS-RELIEF {prompt}, marble texture, cinematic lighting",
+                    prompt=f"BAS-RELIEF {prompt}, marble texture, 8k",
                     image=sr_pil,
                     strength=0.6,
                     num_inference_steps=25,
-                    guidance_scale=7.0
+                    guidance_scale=7.0,
+                    generator=torch.Generator(device).manual_seed(42)
                 ).images[0]
                 logger.success("Estilo aplicado")
             except Exception as e:
                 logger.error(f"Erro no estilo: {str(e)}")
 
         # ========= 3. MAPA DE PROFUNDIDADE =========
-        mapa_pil = arte_pil  # Fallback padrão
-        if modelo_profundidade and arte_pil:
+        mapa_pil = arte_pil  # Fallback
+        if modelo_profundidade:
             try:
                 logger.etapa("Calculando Profundidade")
                 inputs = processador_profundidade(arte_pil, return_tensors="pt").to(device)
@@ -160,13 +158,13 @@ def processar_imagem(imagem, escala, modelo, prompt):
                     depth = modelo_profundidade(**inputs).predicted_depth
 
                 depth = torch.nn.functional.interpolate(
-                    depth.unsqueeze(1).float(),
+                    depth.unsqueeze(1),
                     size=arte_pil.size[::-1],
                     mode="bicubic"
                 ).squeeze().cpu().numpy()
 
-                depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())
-                mapa_pil = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+                depth = (depth - depth.min()) / (depth.max() - depth.min())
+                mapa_pil = Image.fromarray((depth * 255).astype(np.uint8))
                 logger.success("Profundidade calculada")
             except Exception as e:
                 logger.error(f"Erro na profundidade: {str(e)}")
@@ -180,18 +178,16 @@ def processar_imagem(imagem, escala, modelo, prompt):
 
 # ================== INTERFACE GRADIO ==================
 with gr.Blocks(title="TheraSR Universal", theme=gr.themes.Soft()) as app:
-    gr.Markdown("# 🏛 TheraSR - Super Resolução & Arte")
+    gr.Markdown("# 🏛 TheraSR - Processamento Completo em Float32")
 
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(label="Imagem de Entrada", type="pil")
-            scale = gr.Slider(1.0, 4.0, value=2.0,
-                              label="Fator de Escala", step=0.1)
-            model_select = gr.Radio(["EDSR", "RDN"],
-                                    value="EDSR", label="Modelo")
-            style_prompt = gr.Textbox(
-                label="Descrição do Estilo",
-                value="ancient greek marble浮雕, ultra detailed, 8k"
+            input_img = gr.Image(label="Imagem de Entrada", type="pil")
+            slider_scale = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
+            radio_model = gr.Radio(["EDSR", "RDN"], value="EDSR", label="Modelo")
+            text_prompt = gr.Textbox(
+                label="Prompt de Estilo",
+                value="ancient marble浮雕, ultra detailed, 8k cinematic"
             )
             btn_process = gr.Button("Processar", variant="primary")
 
@@ -201,8 +197,8 @@ with gr.Blocks(title="TheraSR Universal", theme=gr.themes.Soft()) as app:
             output_depth = gr.Image(label="Mapa de Profundidade", interactive=False)
 
     btn_process.click(
-        processar_imagem,
-        inputs=[input_image, scale, model_select, style_prompt],
+        processar_imagem_completa,
+        inputs=[input_img, slider_scale, radio_model, text_prompt],
         outputs=[output_sr, output_art, output_depth]
     )