fixing bugs

app.py

@@ -4,25 +4,25 @@ import numpy as np
 import jax
 import pickle
 from PIL import Image
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, file_download
 from model import build_thera
 from super_resolve import process
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline
 from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 
+# Fix de compatibilidade
+file_download.cached_download = file_download.hf_hub_download
+
 # ========== Configuração do Thera ==========
 REPO_ID_EDSR = "prs-eth/thera-edsr-pro"
 REPO_ID_RDN = "prs-eth/thera-rdn-pro"
 
 
-# Carregar modelos Thera
 def load_thera_model(repo_id):
     model_path = hf_hub_download(repo_id=repo_id, filename="model.pkl")
     with open(model_path, 'rb') as fh:
         check = pickle.load(fh)
-        params, backbone, size = check['model'], check['backbone'], check['size']
-        model = build_thera(3, backbone, size)
-    return model, params
+        return build_thera(3, check['backbone'], check['size']), check['model']
 
 
 model_edsr, params_edsr = load_thera_model(REPO_ID_EDSR)
@@ -32,8 +32,7 @@ model_rdn, params_rdn = load_thera_model(REPO_ID_RDN)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Carregar modelos de geração
-pipe = StableDiffusionXLPipeline.from_pretrained(
+pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch_dtype
 ).to(device)
@@ -48,82 +47,66 @@ feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
 depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device)
 
 
-# ========== Funções Principais ==========
-def super_resolution(image, scale_factor, model_type):
-    model = model_edsr if model_type == "EDSR" else model_rdn
-    params = params_edsr if model_type == "EDSR" else params_rdn
-
-    source = np.asarray(image) / 255.
-    target_shape = (
-        round(source.shape[0] * scale_factor),
-        round(source.shape[1] * scale_factor),
-    )
-
-    output = process(source, model, params, target_shape, do_ensemble=True)
-    return Image.fromarray(np.asarray(output))
-
-
-def generate_bas_relief(prompt):
-    full_prompt = f"BAS-RELIEF {prompt}"
-    image = pipe(
-        prompt=full_prompt,
+# ========== Fluxo Integrado ==========
+def full_pipeline(image, scale_factor, model_type, style_prompt):
+    # 1. Super-Resolution
+    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
+    prompt = f"BAS-RELIEF {style_prompt}, intricate carving, marble texture"
+    bas_relief = pipe(
+        prompt=prompt,
+        image=sr_image,
+        strength=0.6,
         num_inference_steps=25,
-        guidance_scale=7.5,
-        height=512,
-        width=512
+        guidance_scale=7.5
     ).images[0]
 
-    inputs = feature_extractor(image, return_tensors="pt").to(device)
+    # 3. Depth Map Estimation
+    inputs = feature_extractor(bas_relief, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = depth_model(**inputs)
-        depth_map = outputs.predicted_depth
+        depth = outputs.predicted_depth
 
-    depth_map = torch.nn.functional.interpolate(
-        depth_map.unsqueeze(1),
-        size=image.size[::-1],
+    depth = torch.nn.functional.interpolate(
+        depth.unsqueeze(1),
+        size=bas_relief.size[::-1],
         mode="bicubic"
     ).squeeze().cpu().numpy()
 
-    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
-    depth_map = (depth_map * 255).astype(np.uint8)
+    depth = (depth - depth.min()) / (depth.max() - depth.min())
+    depth = (depth * 255).astype(np.uint8)
 
-    return image, Image.fromarray(depth_map)
+    return sr_image, bas_relief, Image.fromarray(depth)
 
 
 # ========== Interface Gradio ==========
-with gr.Blocks(title="TheraSR + Bas-Relief Generator") as app:
-    gr.Markdown("# 🔥 TheraSR + Bas-Relief Generator")
-    gr.Markdown("Combine aliasing-free super-resolution with artistic bas-relief generation")
-
-    with gr.Tabs():
-        with gr.TabItem("🖼 Super-Resolution"):
-            with gr.Row():
-                sr_input = gr.Image(label="Input Image", type="pil")
-                sr_output = gr.Image(label="Super-Resolution Result")
-            sr_scale = gr.Slider(1.0, 6.0, value=2.0, label="Scale Factor")
-            sr_model = gr.Radio(["EDSR", "RDN"], value="EDSR", label="Model Type")
-            sr_btn = gr.Button("Enhance Resolution")
-
-        with gr.TabItem("🎨 Generate Bas-Relief"):
-            with gr.Row():
-                text_input = gr.Textbox(label="Art Prompt", placeholder="Roman soldier marble relief...")
-            with gr.Row():
-                gen_output = gr.Image(label="Generated Art")
-                depth_output = gr.Image(label="Depth Map")
-            gen_btn = gr.Button("Generate Artwork")
-
-    # Event Handlers
-    sr_btn.click(
-        super_resolution,
-        inputs=[sr_input, sr_scale, sr_model],
-        outputs=sr_output
-    )
-
-    gen_btn.click(
-        generate_bas_relief,
-        inputs=text_input,
-        outputs=[gen_output, depth_output]
+with gr.Blocks(title="TheraSR + Bas-Relief Fusion") as app:
+    gr.Markdown("## 🪄 Super-Resolution → Bas-Relief → Depth Map")
+
+    with gr.Row():
+        with gr.Column():
+            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")
+            process_btn = gr.Button("Start Pipeline")
+
+        with gr.Column():
+            sr_output = gr.Image(label="Super-Resolution Result")
+            style_output = gr.Image(label="Bas-Relief Result")
+            depth_output = gr.Image(label="Depth Map")
+
+    process_btn.click(
+        full_pipeline,
+        inputs=[input_image, scale, model_type, style_prompt],
+        outputs=[sr_output, style_output, depth_output]
     )
 
-# Configuração do Hugging Face
-app.launch(debug=False, share=True)
+app.launch(debug=False)