import torch
import gradio as gr
from diffusers import ShapEPipeline
from diffusers.utils import export_to_gif
import os
from huggingface_hub import HfApi, login
import json

# GPU kontrolü
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Kullanılan cihaz: {device}")

def validate_token(token):
    try:
        login(token=token)
        return True
    except Exception as e:
        print(f"Token doğrulama hatası: {str(e)}")
        return False

def generate_3d_model(prompt, token, guidance_scale=15.0, num_steps=64, export_format="obj"):
    try:
        if not validate_token(token):
            return "Geçersiz Hugging Face token'ı", None, None
            
        print(f"Üretim başlıyor: {prompt}")
        
        # Model yükleme (token ile)
        pipe = ShapEPipeline.from_pretrained(
            "openai/shap-e",
            torch_dtype=torch.float32 if device == "cuda" else torch.float32,
            token=token,
            cache_dir="model_cache"
        ).to(device)
        
        # Çıktı klasörü kontrolü
        os.makedirs("outputs", exist_ok=True)
        
        # Dosya adı oluşturma
        safe_prompt = "".join(x for x in prompt if x.isalnum() or x in (" ", "-", "_"))
        base_filename = f"outputs/{safe_prompt}"
        
        try:
            # 3D model oluşturma
            output = pipe(
                prompt_embeds=None,
                text=prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_steps
            )
            
            images = output.images
            
            # GIF kaydetme
            gif_path = export_to_gif(images, f"{base_filename}.gif")
            
            # 3D mesh oluşturma
            mesh_output = pipe(
                prompt_embeds=None,
                text=prompt,
                guidance_scale=guidance_scale,
                num_inference_steps=num_steps,
                output_type="mesh"
            )
            
            mesh = mesh_output.meshes[0]
            
            # Mesh'i kaydetme
            output_path = f"{base_filename}.{export_format}"
            mesh.export(output_path)
            
            print(f"Üretim tamamlandı: {output_path}")
            return "Üretim başarılı!", gif_path, output_path
            
        except Exception as model_error:
            error_msg = f"Model çalıştırma hatası: {str(model_error)}"
            print(error_msg)
            return error_msg, None, None
        
    except Exception as e:
        error_msg = f"Genel hata: {str(e)}"
        print(error_msg)
        return error_msg, None, None

# Gradio arayüzü
interface = gr.Interface(
    fn=generate_3d_model,
    inputs=[
        gr.Textbox(label="3D Model için açıklama giriniz (İngilizce)"),
        gr.Textbox(label="Hugging Face Token", type="password"),
        gr.Slider(minimum=1, maximum=20, value=15, label="Guidance Scale"),
        gr.Slider(minimum=32, maximum=128, value=64, label="Adım Sayısı"),
        gr.Radio(["obj", "glb"], label="Export Formatı", value="obj")
    ],
    outputs=[
        gr.Textbox(label="Durum"),
        gr.Image(label="3D Önizleme (GIF)"),
        gr.File(label="3D Model Dosyası")
    ],
    title="DNM3D - 3D Model Üretici",
    description="Metin girişiyle 3D model oluşturun. Kullanmak için Hugging Face token'ınızı girmeniz gerekiyor."
)

if __name__ == "__main__":
    interface.launch()