fixing bugs

app.py

@@ -23,19 +23,14 @@ class CustomLogger:
         self.logger.addHandler(handler)
         self.logger.setLevel(logging.INFO)
 
-
     def divider(self, text=None, length=60):
         if text:
-            # Cálculo seguro do número de '='
-            available_space = length - len(text) - 12  # 10 '=' + 2 espaços
-            if available_space < 1:
-                available_space = 1  # Garante pelo menos 1 '='
-            msg = f"\n{'=' * 10} {text.upper()} {'=' * available_space}"
+            available_space = max(length - len(text) - 12, 1)
+            msg = f"{'=' * 10} {text.upper()} {'=' * available_space}"
         else:
-            msg = "\n" + "=" * length
+            msg = "=" * length
         self.logger.info(msg)
 
-
     def etapa(self, text):
         self.logger.info(f"▶ {text}")
 
@@ -53,118 +48,116 @@ logger = CustomLogger(__name__)
 
 # ================== CONFIGURAÇÃO DE HARDWARE ==================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if device == "cuda" else torch.float32
-logger.divider("Configuração Inicial")
+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"Precisão numérica: {str(torch_dtype).replace('torch.', '')}")
+logger.success(f"Modo de precisão: float32")
 
 
 # ================== CARREGAMENTO DE MODELOS ==================
 def carregar_modelo_thera(repo_id):
+    """Carrega modelos Thera com tratamento de erros robusto"""
     try:
-        logger.divider(f"Carregando Modelo: {repo_id}")
+        logger.divider(f"Carregando {repo_id}")
         model_path = hf_hub_download(repo_id=repo_id, filename="model.pkl")
         with open(model_path, 'rb') as f:
             check = pickle.load(f)
             model = build_thera(3, check['backbone'], check['size'])
             params = check['model']
-        logger.success(f"Modelo {repo_id} carregado")
+        logger.success(f"{repo_id} carregado")
         return model, params
     except Exception as e:
         logger.error(f"Falha ao carregar {repo_id}: {str(e)}")
         return None, None
 
 
-# Carregar modelos Thera
-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 no carregamento dos modelos Thera")
-    raise
+# 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")
 
-# ================== PIPELINE DE ARTE ==================
+# ================== MODELOS DE ARTE (CARREGAMENTO CONDICIONAL) ==================
 pipe = None
 modelo_profundidade = None
+processador_profundidade = None
 
 try:
-    logger.divider("Configurando Pipeline de Arte")
+    logger.divider("Inicializando Componentes Artísticos")
+
+    # Pipeline de estilo
     pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0",
         torch_dtype=torch_dtype,
-        variant="fp16",
         use_safetensors=True
     ).to(device)
 
+    # Adapter LoRA
     pipe.load_lora_weights(
         "KappaNeuro/bas-relief",
         weight_name="BAS-RELIEF.safetensors"
     )
-    logger.success("Pipeline SDXL configurado")
 
-    logger.etapa("Configurando Modelo de Profundidade")
+    # Modelo de profundidade
     processador_profundidade = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
     modelo_profundidade = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
-    logger.success("Modelo de profundidade pronto")
 
+    logger.success("Componentes artísticos prontos")
 except Exception as e:
-    logger.error(f"Erro na configuração da GPU: {str(e)}")
+    logger.warning(f"Recursos artísticos desativados: {str(e)}")
     pipe = None
     modelo_profundidade = None
 
 
-# ================== FLUXO DE PROCESSAMENTO ==================
-def pipeline_completo(imagem, fator_escala, modelo_escolhido, prompt_estilo):
+# ================== PIPELINE PRINCIPAL ==================
+def processar_imagem(imagem, escala, modelo, prompt):
+    """Fluxo completo de processamento com fallbacks"""
     try:
         logger.divider("Novo Processamento")
 
-        # Converter entrada
+        # Converter entrada para PIL
         if not isinstance(imagem, Image.Image):
             imagem = Image.fromarray(imagem)
 
-        # ========= SUPER-RESOLUÇÃO =========
-        logger.etapa("Processando Super-Resolução")
-        modelo = modelo_edsr if modelo_escolhido == "EDSR" else modelo_rdn
-        params = params_edsr if modelo_escolhido == "EDSR" else params_rdn
-
-        sr_array = process(
-            np.array(imagem) / 255.,
-            modelo,
-            params,
-            (round(imagem.size[1] * fator_escala),
-             round(imagem.size[0] * fator_escala)),
+        # ========= 1. SUPER-RESOLUÇÃO =========
+        logger.etapa("Super-Resolução Thera")
+        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,
+            modelo_sr,
+            params_sr,
+            (int(imagem.height * escala),
+             int(imagem.width * escala)),
             True
         )
 
-        sr_pil = Image.fromarray(np.array(sr_array)).convert("RGB")
-        logger.success(f"Super-Resolução: {sr_pil.size[0]}x{sr_pil.size[1]}")
+        sr_pil = Image.fromarray(np.array(sr_jax)).convert("RGB")
+        logger.success(f"Resolução: {sr_pil.size[0]}x{sr_pil.size[1]}")
 
-        # ========= ESTILO BAIXO-RELEVO =========
-        arte_pil = None
-        if pipe and modelo_profundidade:
+        # ========= 2. ESTILO BAIXO-RELEVO =========
+        arte_pil = sr_pil  # Fallback padrão
+        if pipe:
             try:
-                logger.etapa("Aplicando Estilo Artístico")
-                resultado = pipe(
-                    prompt=f"BAS-RELIEF {prompt_estilo}, intricate marble carving, 8k ultra HD",
+                logger.etapa("Aplicando Estilo")
+                arte_pil = pipe(
+                    prompt=f"BAS-RELIEF {prompt}, marble texture, cinematic lighting",
                     image=sr_pil,
-                    strength=0.65,
-                    num_inference_steps=30,
-                    guidance_scale=7.5
-                )
-                arte_pil = resultado.images[0]
-                logger.success(f"Arte gerada: {arte_pil.size[0]}x{arte_pil.size[1]}")
+                    strength=0.6,
+                    num_inference_steps=25,
+                    guidance_scale=7.0
+                ).images[0]
+                logger.success("Estilo aplicado")
             except Exception as e:
-                logger.error(f"Falha no estilo: {str(e)}")
+                logger.error(f"Erro no estilo: {str(e)}")
 
-        # ========= MAPA DE PROFUNDIDADE =========
-        mapa_pil = None
-        if arte_pil and modelo_profundidade:
+        # ========= 3. MAPA DE PROFUNDIDADE =========
+        mapa_pil = arte_pil  # Fallback padrão
+        if modelo_profundidade and arte_pil:
             try:
                 logger.etapa("Calculando Profundidade")
-                inputs = processador_profundidade(images=arte_pil, return_tensors="pt").to(device)
+                inputs = processador_profundidade(arte_pil, return_tensors="pt").to(device)
                 with torch.no_grad():
-                    outputs = modelo_profundidade(**inputs)
-                    depth = outputs.predicted_depth
+                    depth = modelo_profundidade(**inputs).predicted_depth
 
                 depth = torch.nn.functional.interpolate(
                     depth.unsqueeze(1).float(),
@@ -172,47 +165,45 @@ def pipeline_completo(imagem, fator_escala, modelo_escolhido, prompt_estilo):
                     mode="bicubic"
                 ).squeeze().cpu().numpy()
 
-                depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
-                mapa_pil = Image.fromarray((depth * 255).astype(np.uint8))
-                logger.success("Mapa de profundidade calculado")
+                depth_normalized = (depth - depth.min()) / (depth.max() - depth.min())
+                mapa_pil = Image.fromarray((depth_normalized * 255).astype(np.uint8))
+                logger.success("Profundidade calculada")
             except Exception as e:
-                logger.error(f"Falha na profundidade: {str(e)}")
+                logger.error(f"Erro na profundidade: {str(e)}")
 
-        return sr_pil, arte_pil or sr_pil, mapa_pil or sr_pil
+        return sr_pil, arte_pil, mapa_pil
 
     except Exception as e:
-        logger.error(f"Erro no pipeline: {str(e)}")
+        logger.error(f"Erro fatal: {str(e)}")
         return None, None, None
 
 
 # ================== INTERFACE GRADIO ==================
-with gr.Blocks(title="TheraSR Art Suite", theme=gr.themes.Soft()) as app:
-    gr.Markdown("# 🎨 TheraSR - Super Resolução & Arte Generativa")
+with gr.Blocks(title="TheraSR Universal", theme=gr.themes.Soft()) as app:
+    gr.Markdown("# 🏛 TheraSR - Super Resolução & Arte")
 
     with gr.Row():
         with gr.Column():
-            entrada_imagem = gr.Image(label="Imagem de Entrada", type="pil")
-            seletor_modelo = gr.Radio(
-                ["EDSR", "RDN"],
-                value="EDSR",
-                label="Modelo de Super-Resolução"
-            )
-            controle_escala = gr.Slider(1.0, 4.0, value=2.0, label="Fator de Escala")
-            entrada_prompt = gr.Textbox(
-                label="Prompt de Estilo",
-                value="insanely detailed ancient greek marble浮雕, 8k cinematic lighting"
+            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"
             )
-            botao_processar = gr.Button("Gerar", variant="primary")
+            btn_process = gr.Button("Processar", variant="primary")
 
         with gr.Column():
-            saida_sr = gr.Image(label="Super-Resolução", show_label=True)
-            saida_arte = gr.Image(label="Arte em Relevo", show_label=True)
-            saida_profundidade = gr.Image(label="Mapa de Profundidade", show_label=True)
-
-    botao_processar.click(
-        pipeline_completo,
-        inputs=[entrada_imagem, controle_escala, seletor_modelo, entrada_prompt],
-        outputs=[saida_sr, saida_arte, saida_profundidade]
+            output_sr = gr.Image(label="Super-Resolução", interactive=False)
+            output_art = gr.Image(label="Arte em Relevo", interactive=False)
+            output_depth = gr.Image(label="Mapa de Profundidade", interactive=False)
+
+    btn_process.click(
+        processar_imagem,
+        inputs=[input_image, scale, model_select, style_prompt],
+        outputs=[output_sr, output_art, output_depth]
     )
 
 if __name__ == "__main__":