# import os
# import subprocess

# # Убираем pyenv, если вдруг остался .python-version
# os.environ.pop("PYENV_VERSION", None)

# # Установка зависимостей
# subprocess.run(["pip", "install", "torch", "wheel"], check=True)
# subprocess.run([
#     "pip", "install", "--no-build-isolation", 
#     "diso@git+https://github.com/SarahWeiii/diso.git"
# ], check=True)

# # Импорты
# import gradio as gr
# import uuid
# import torch
# import zipfile
# import requests
# import traceback
# import trimesh
# from trimesh.exchange.gltf import export_glb 

# print("Trimesh version:", trimesh.__version__)


# from inference_triposg import run_triposg
# from triposg.pipelines.pipeline_triposg import TripoSGPipeline
# from briarmbg import BriaRMBG

# # Настройки устройства
# device = "cuda" if torch.cuda.is_available() else "cpu"
# dtype = torch.float16 if device == "cuda" else torch.float32

# # Загрузка весов
# weights_dir = "pretrained_weights"
# triposg_path = os.path.join(weights_dir, "TripoSG")
# rmbg_path = os.path.join(weights_dir, "RMBG-1.4")

# if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
#     print("📦 Downloading pretrained weights...")
#     url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
#     zip_path = "pretrained_models.zip"

#     with requests.get(url, stream=True) as r:
#         r.raise_for_status()
#         with open(zip_path, "wb") as f:
#             for chunk in r.iter_content(chunk_size=8192):
#                 f.write(chunk)

#     print("📦 Extracting weights...")
#     with zipfile.ZipFile(zip_path, "r") as zip_ref:
#         zip_ref.extractall(weights_dir)

#     os.remove(zip_path)
#     print("✅ Weights ready.")

# # Загрузка моделей
# pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
# rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
# rmbg_net.eval()

# # Генерация .glb
# # def generate(image_path):
# def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
#     print("[API CALL] image_path received:", image_path)
#     print("[API CALL] File exists:", os.path.exists(image_path))

#     temp_id = str(uuid.uuid4())
#     output_path = f"/tmp/{temp_id}.glb"

#     print("[DEBUG] Generating mesh from:", image_path)

#     try:
#         # mesh = run_triposg(
#         #     pipe=pipe,
#         #     image_input=image_path,
#         #     rmbg_net=rmbg_net,
#         #     seed=42,
#         #     num_inference_steps=25,
#         #     guidance_scale=5.0,
#         #     faces=-1,
#         # )
#         mesh = run_triposg(
#             pipe=pipe,
#             image_input=image_path,
#             rmbg_net=rmbg_net,
#             seed=42,
#             num_inference_steps=int(num_steps),
#             guidance_scale=float(guidance_scale),
#             faces=int(face_number),
#         )

#         # if mesh is None:
#         #     raise ValueError("Mesh generation failed")

#         # mesh.export(output_path)
#         # print(f"[DEBUG] Mesh saved to {output_path}")

#         # return output_path if os.path.exists(output_path) else "Error: output file not found"

        
#         # if mesh is None:
#         #     raise ValueError("Mesh generation failed")

#         # # Убираем визуал, метаданные, обертки
#         # mesh.visual = None
#         # mesh.metadata.clear()
#         # mesh.name = "endless_tools"

#         # # Экспорт только геометрии
#         # glb_data = mesh.export(file_type="glb")
#         # with open(output_path, "wb") as f:
#         #     f.write(glb_data)

#         # print(f"[DEBUG] Mesh saved to {output_path}")

#         # return output_path if os.path.exists(output_path) else "Error: output file not found"

#         if mesh is None:
#             raise ValueError("Mesh generation returned None")

#         # Очистка визуала, метаданных и имени
#         mesh.visual = None
#         mesh.metadata.clear()
#         mesh.name = "geometry_0"

#         # glb_data = mesh.export(file_type="glb")
#         glb_data = export_glb(mesh)
#         with open(output_path, "wb") as f:
#             f.write(glb_data)

#         # Экспорт .glb вручную (иначе Trimesh добавляет сцену)
#         # glb_data = mesh.export(file_type="glb")
#         # with open(output_path, "wb") as f:
#         #     f.write(glb_data)
        

#         print(f"[DEBUG] Mesh saved to {output_path}")
#         return output_path if os.path.exists(output_path) else None
#     # except Exception as e:
#     #     print("[ERROR]", e)
#     #     return f"Error: {e}"
#     except Exception as e:
#         # import traceback
#         print("[ERROR]", e)
#         traceback.print_exc()  # ← выведет полную трассировку в логи
#         return f"Error: {e}"

# # Интерфейс Gradio
# demo = gr.Interface(
#     fn=generate,
#     inputs=gr.Image(type="filepath", label="Upload image"),
#     outputs=gr.File(label="Download .glb"),
#     title="TripoSG Image to 3D",
#     description="Upload an image to generate a 3D model (.glb)",
# )

# # Запуск
# demo.launch()













# import os
# import subprocess

# # Убираем pyenv
# os.environ.pop("PYENV_VERSION", None)

# # Установка зависимостей
# subprocess.run(["pip", "install", "torch", "wheel"], check=True)
# subprocess.run([
#     "pip", "install", "--no-build-isolation", 
#     "diso@git+https://github.com/SarahWeiii/diso.git"
# ], check=True)

# # Импорты (перенесены после установки зависимостей)
# import gradio as gr
# import uuid
# import torch
# import zipfile
# import requests
# import traceback
# import trimesh
# from trimesh.exchange.gltf import export_glb

# from inference_triposg import run_triposg
# from triposg.pipelines.pipeline_triposg import TripoSGPipeline
# from briarmbg import BriaRMBG


# print("Trimesh version:", trimesh.__version__)

# # Настройки устройства
# device = "cuda" if torch.cuda.is_available() else "cpu"
# dtype = torch.float16 if device == "cuda" else torch.float32

# # Загрузка весов
# weights_dir = "pretrained_weights"
# triposg_path = os.path.join(weights_dir, "TripoSG")
# rmbg_path = os.path.join(weights_dir, "RMBG-1.4")

# if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
#     print("📦 Downloading pretrained weights...")
#     url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
#     zip_path = "pretrained_models.zip"

#     with requests.get(url, stream=True) as r:
#         r.raise_for_status()
#         with open(zip_path, "wb") as f:
#             for chunk in r.iter_content(chunk_size=8192):
#                 f.write(chunk)

#     print("📦 Extracting weights...")
#     with zipfile.ZipFile(zip_path, "r") as zip_ref:
#         zip_ref.extractall(weights_dir)

#     os.remove(zip_path)
#     print("✅ Weights ready.")

# # Загрузка моделей
# pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
# rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
# rmbg_net.eval()

# # Генерация .glb
# def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
#     print("[API CALL] image_path received:", image_path)
#     print("[API CALL] File exists:", os.path.exists(image_path))

#     temp_id = str(uuid.uuid4())
#     output_path = f"/tmp/{temp_id}.glb"
#     print("[DEBUG] Generating mesh from:", image_path)

#     try:
#         mesh = run_triposg(
#             pipe=pipe,
#             image_input=image_path,
#             rmbg_net=rmbg_net,
#             seed=42,
#             num_inference_steps=int(num_steps),
#             guidance_scale=float(guidance_scale),
#             faces=int(face_number),
#         )

#         if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
#             raise ValueError("Mesh generation returned an empty mesh")

#         # Безопасная очистка визуала
#         if hasattr(mesh, "visual") and mesh.visual is not None:
#             try:
#                 mesh.visual = None
#             except Exception:
#                 print("[WARN] Failed to clear visual, skipping")

#         mesh.metadata.clear()
#         mesh.name = "endless_tools_mesh"

#         # Экспорт .glb
#         # glb_data = export_glb(mesh)
#         glb_data = mesh.export(file_type='glb')
#         with open(output_path, "wb") as f:
#             f.write(glb_data)

#         print(f"[DEBUG] Mesh saved to {output_path}")
#         return output_path if os.path.exists(output_path) else None

#     except Exception as e:
#         print("[ERROR]", e)
#         traceback.print_exc()
#         return f"Error: {e}"

# # Интерфейс Gradio
# demo = gr.Interface(
#     fn=generate,
#     inputs=gr.Image(type="filepath", label="Upload image"),
#     outputs=gr.File(label="Download .glb"),
#     title="TripoSG Image to 3D",
#     description="Upload an image to generate a 3D model (.glb)",
# )

# # Запуск
# demo.launch()









# import os
# import subprocess

# # Убираем pyenv
# os.environ.pop("PYENV_VERSION", None)

# # Установка зависимостей
# subprocess.run(["pip", "install", "torch", "wheel"], check=True)
# subprocess.run([
#     "pip", "install", "--no-build-isolation", 
#     "diso@git+https://github.com/SarahWeiii/diso.git"
# ], check=True)

# # Импорты (перенесены после установки зависимостей)
# import gradio as gr
# import uuid
# import torch
# import zipfile
# import requests
# import traceback
# import trimesh
# import numpy as np
# from trimesh.exchange.gltf import export_glb

# from inference_triposg import run_triposg
# from triposg.pipelines.pipeline_triposg import TripoSGPipeline
# from briarmbg import BriaRMBG


# print("Trimesh version:", trimesh.__version__)

# # Настройки устройства
# device = "cuda" if torch.cuda.is_available() else "cpu"
# dtype = torch.float16 if device == "cuda" else torch.float32

# # Загрузка весов
# weights_dir = "pretrained_weights"
# triposg_path = os.path.join(weights_dir, "TripoSG")
# rmbg_path = os.path.join(weights_dir, "RMBG-1.4")

# if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
#     print("📦 Downloading pretrained weights...")
#     url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
#     zip_path = "pretrained_models.zip"

#     with requests.get(url, stream=True) as r:
#         r.raise_for_status()
#         with open(zip_path, "wb") as f:
#             for chunk in r.iter_content(chunk_size=8192):
#                 f.write(chunk)

#     print("📦 Extracting weights...")
#     with zipfile.ZipFile(zip_path, "r") as zip_ref:
#         zip_ref.extractall(weights_dir)

#     os.remove(zip_path)
#     print("✅ Weights ready.")

# # Загрузка моделей
# pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
# rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
# rmbg_net.eval()

# # Генерация .glb
# # def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
# #     print("[API CALL] image_path received:", image_path)
# #     print("[API CALL] File exists:", os.path.exists(image_path))

# #     temp_id = str(uuid.uuid4())
# #     output_path = f"/tmp/{temp_id}.glb"
# #     print("[DEBUG] Generating mesh from:", image_path)

# #     try:
# #         mesh = run_triposg(
# #             pipe=pipe,
# #             image_input=image_path,
# #             rmbg_net=rmbg_net,
# #             seed=42,
# #             num_inference_steps=int(num_steps),
# #             guidance_scale=float(guidance_scale),
# #             faces=int(face_number),
# #         )

# #         if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
# #             raise ValueError("Mesh generation returned an empty mesh")

# #         mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
# #         mesh.rezero()
# #         mesh.fix_normals()
# #         mesh.apply_translation(-mesh.center_mass)  

# #         # Масштабируем, чтобы модель вписывалась в размер 1x1x1 
# #         # Если нужно будет подгонять под размер в Endless Tools, то можно использовать:
# #         # scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
# #         # mesh.apply_scale(scale_factor)


# #         glb_data = mesh.export(file_type='glb')
# #         with open(output_path, "wb") as f:
# #             f.write(glb_data)

# #         print(f"[DEBUG] Mesh saved to {output_path}")
# #         return output_path if os.path.exists(output_path) else None

# #     except Exception as e:
# #         print("[ERROR]", e)
# #         traceback.print_exc()
# #         return f"Error: {e}"

# def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
#     print("[API CALL] image_path received:", image_path)
#     print("[API CALL] File exists:", os.path.exists(image_path))

#     temp_id = str(uuid.uuid4())
#     output_path = f"/tmp/{temp_id}.glb"
#     print("[DEBUG] Generating mesh from:", image_path)

#     try:
#         mesh = run_triposg(
#             pipe=pipe,
#             image_input=image_path,
#             rmbg_net=rmbg_net,
#             seed=42,
#             num_inference_steps=int(num_steps),
#             guidance_scale=float(guidance_scale),
#             faces=int(face_number),
#         )

#         if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
#             raise ValueError("Mesh generation returned an empty mesh")

#         # 🔧 Пересоздаём Trimesh и гарантируем чистоту геометрии
#         mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, process=True)

#         # ✅ Центрируем модель
#         mesh.apply_translation(-mesh.center_mass)

#         # ✅ Масштабируем к единичному размеру (все модели ~одинаковые)
#         scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
#         mesh.apply_scale(scale_factor)

#         # ✅ Гарантированно пересчитываем нормали
#         mesh.fix_normals()

#         # print("[DEBUG] Normals present:", mesh.has_vertex_normals)
#         if hasattr(mesh, "vertex_normals"):
#             print("[DEBUG] Normals shape:", mesh.vertex_normals.shape)
#         else:
#             print("[DEBUG] Normals missing.")


#         # 💾 Сохраняем GLB
#         glb_data = mesh.export(file_type='glb')
#         with open(output_path, "wb") as f:
#             f.write(glb_data)

#         print(f"[DEBUG] Mesh saved to {output_path}")
#         return output_path if os.path.exists(output_path) else None

#     except Exception as e:
#         print("[ERROR]", e)
#         traceback.print_exc()
#         return f"Error: {e}"


# # Интерфейс Gradio
# demo = gr.Interface(
#     fn=generate,
#     inputs=gr.Image(type="filepath", label="Upload image"),
#     outputs=gr.File(label="Download .glb"),
#     title="TripoSG Image to 3D",
#     description="Upload an image to generate a 3D model (.glb)",
# )

# # Запуск
# demo.launch()










import os
import subprocess

# Убираем pyenv
os.environ.pop("PYENV_VERSION", None)

# Установка зависимостей
subprocess.run(["pip", "install", "torch", "wheel"], check=True)
subprocess.run([
    "pip", "install", "--no-build-isolation", 
    "diso@git+https://github.com/SarahWeiii/diso.git"
], check=True)

# Импорты (перенесены после установки зависимостей)
import gradio as gr
import uuid
import torch
import zipfile
import requests
import traceback
import trimesh
import numpy as np


from trimesh.exchange.gltf import export_glb

from inference_triposg import run_triposg
from triposg.pipelines.pipeline_triposg import TripoSGPipeline
from briarmbg import BriaRMBG

GLTF_PACK = "/tmp/gltfpack"

print("Trimesh version:", trimesh.__version__)

# Настройки устройства
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32

# Загрузка весов
weights_dir = "pretrained_weights"
triposg_path = os.path.join(weights_dir, "TripoSG")
rmbg_path = os.path.join(weights_dir, "RMBG-1.4")

if not (os.path.exists(triposg_path) and os.path.exists(rmbg_path)):
    print("📦 Downloading pretrained weights...")
    url = "https://huggingface.co/datasets/endlesstools/pretrained-assets/resolve/main/pretrained_models.zip"
    zip_path = "pretrained_models.zip"

    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(zip_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)

    print("📦 Extracting weights...")
    with zipfile.ZipFile(zip_path, "r") as zip_ref:
        zip_ref.extractall(weights_dir)

    os.remove(zip_path)
    print("✅ Weights ready.")

# Загрузка моделей
pipe = TripoSGPipeline.from_pretrained(triposg_path).to(device, dtype)
rmbg_net = BriaRMBG.from_pretrained(rmbg_path).to(device)
rmbg_net.eval()


def generate(image_path, face_number=50000, guidance_scale=5.0, num_steps=25):
    print("[API CALL] image_path received:", image_path)
    print("[API CALL] File exists:", os.path.exists(image_path))

    temp_id = str(uuid.uuid4())
    output_path = f"/tmp/{temp_id}.glb"
    print("[DEBUG] Generating mesh from:", image_path)

    try:
        mesh = run_triposg(
            pipe=pipe,
            image_input=image_path,
            rmbg_net=rmbg_net,
            seed=42,
            num_inference_steps=int(num_steps),
            guidance_scale=float(guidance_scale),
            faces=int(face_number),
        )

        if mesh is None or mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
            raise ValueError("Mesh generation returned an empty mesh")

        # 🔧 Пересоздаём Trimesh и гарантируем чистоту геометрии
        mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, process=True)

        # ✅ Центрируем модель
        mesh.apply_translation(-mesh.center_mass)

        # ✅ Масштабируем к единичному размеру (все модели ~одинаковые)
        scale_factor = 1.0 / np.max(np.linalg.norm(mesh.vertices, axis=1))
        mesh.apply_scale(scale_factor)

        # ✅ Гарантированно пересчитываем нормали
        mesh.fix_normals()

        # print("[DEBUG] Normals present:", mesh.has_vertex_normals)
        if hasattr(mesh, "vertex_normals"):
            print("[DEBUG] Normals shape:", mesh.vertex_normals.shape)
        else:
            print("[DEBUG] Normals missing.")

        # 💾 Сохраняем GLB
        glb_data = mesh.export(file_type='glb')
        with open(output_path, "wb") as f:
            f.write(glb_data)

        print(f"[DEBUG] Mesh saved to {output_path}")
        return output_path if os.path.exists(output_path) else None

    except Exception as e:
        print("[ERROR]", e)
        traceback.print_exc()
        return f"Error: {e}"


# Интерфейс Gradio
demo = gr.Interface(
    fn=generate,
    inputs=gr.Image(type="filepath", label="Upload image"),
    outputs=gr.File(label="Download .glb"),
    title="TripoSG Image to 3D",
    description="Upload an image to generate a 3D model (.glb)",
)



# Запуск
demo.launch()