fixing bugs

app.py

@@ -46,47 +46,55 @@ pipe.load_lora_weights(
     peft_backend="peft"
 )
 
-feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
+# ========== Configuração do Modelo de Profundidade ==========
+depth_processor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")  # Nome padronizado
 depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
 
-
 # ========== Fluxo Integrado ==========
 def full_pipeline(image, scale_factor, model_type, style_prompt):
-    # 1. Super-Resolution
+    # 1. Super-Resolution (JAX)
     sr_model = model_edsr if model_type == "EDSR" else model_rdn
     sr_params = params_edsr if model_type == "EDSR" else params_rdn
-    sr_image = process(np.array(image) / 255., sr_model, sr_params,
-                       (round(image.size[1] * scale_factor),
-                        round(image.size[0] * scale_factor)),
-                       True)
 
-    # 2. Bas-Relief Style Transfer
+    # Processar e converter para numpy array
+    sr_jax = process(np.array(image) / 255., sr_model, sr_params,
+                     (round(image.size[1] * scale_factor),
+                      round(image.size[0] * scale_factor)),
+                     True)
+
+    # Conversão crítica: JAX Array → numpy → PIL
+    sr_np = np.asarray(sr_jax)
+    sr_pil = Image.fromarray(sr_np)
+
+    if device == "cpu":
+        return sr_pil, None, None
+
+    # 2. Style Transfer (PyTorch)
     prompt = f"BAS-RELIEF {style_prompt}, intricate carving, marble texture"
     bas_relief = pipe(
         prompt=prompt,
-        image=sr_image,
+        image=sr_pil,  # Usar PIL Image diretamente
         strength=0.6,
         num_inference_steps=25,
         guidance_scale=7.5
     ).images[0]
 
-    # 3. Depth Map Estimation
-    inputs = feature_extractor(bas_relief, return_tensors="pt").to(device)
+    # 3. Depth Map
+    inputs = depth_processor(bas_relief, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = depth_model(**inputs)
         depth = outputs.predicted_depth
 
     depth = torch.nn.functional.interpolate(
         depth.unsqueeze(1),
-        size=bas_relief.size[::-1],
-        mode="bicubic"
+        mode="bicubic",
+        size=bas_relief.size[::-1]
     ).squeeze().cpu().numpy()
 
-    depth = (depth - depth.min()) / (depth.max() - depth.min())
+    depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
     depth = (depth * 255).astype(np.uint8)
 
-    return sr_image, bas_relief, Image.fromarray(depth)
-
+    return sr_pil, bas_relief, Image.fromarray(depth)
 
 # ========== Interface Gradio ==========
 with gr.Blocks(title="TheraSR + Bas-Relief Fusion") as app:
@@ -97,8 +105,10 @@ with gr.Blocks(title="TheraSR + Bas-Relief Fusion") as app:
             input_image = gr.Image(label="Input Image", type="pil")
             scale = gr.Slider(1.0, 4.0, value=2.0, label="Scale Factor")
             model_type = gr.Radio(["EDSR", "RDN"], value="EDSR", label="SR Model")
-            style_prompt = gr.Textbox(label="Style Prompt",
-                                      placeholder="marble sculpture, ancient greek style")
+            style_prompt = gr.Textbox(
+                label="Style Prompt",
+                value="insanely detailed and complex engraving relief, ultra-high definition"  # <-- Alteração aqui
+            )
             process_btn = gr.Button("Start Pipeline")
 
         with gr.Column():