import spaces
import os
import torch
from huggingface_hub import HfApi
from diffusers import WanTransformer3DModel

# -----------------------------------------------------------------------------
# ENV
# -----------------------------------------------------------------------------
os.environ["HF_HOME"] = "/tmp/hf"
os.environ["HF_HUB_CACHE"] = "/tmp/hf/hub"

HF_TOKEN = os.getenv("HF_TOKEN")

MODEL_ID = "ibyteohdear/Wan2.2-I2V-14B-Lightning"
REPO_ID  = "ibyteohdear/Wan2.2-I2V-14B-Lightning-Safe"
SUBFOLDER = "transformer" #_2
DTYPE = torch.bfloat16
device = "cuda"

api = HfApi()

# -----------------------------------------------------------------------------
# CREATE REPO
# -----------------------------------------------------------------------------
api.create_repo(
    repo_id=REPO_ID,
    private=False,
    exist_ok=True,
    token=HF_TOKEN,
)

spaces.GPU(duration=180)
def export():
    # -----------------------------------------------------------------------------
    # LOAD MODEL
    # -----------------------------------------------------------------------------
    print("📥 Loading transformer...")
    model = WanTransformer3DModel.from_pretrained(
        MODEL_ID,
        subfolder=SUBFOLDER,
        torch_dtype=DTYPE,
        low_cpu_mem_usage=True,
        device_map=device,
    ).eval()
    
    # -----------------------------------------------------------------------------
    # SAVE SAFETENSORS
    # -----------------------------------------------------------------------------
    save_dir = "/tmp/wan_transformer_safetensors"
    os.makedirs(save_dir, exist_ok=True)
    
    print("💾 Saving safetensors...")
    model.save_pretrained(
        save_dir,
        safe_serialization=True,
    )
    
    # -----------------------------------------------------------------------------
    # UPLOAD FOLDER (CORRECT WAY)
    # -----------------------------------------------------------------------------
    print("☁️ Uploading to Hugging Face...")
    api.upload_folder(
        folder_path=save_dir,
        repo_id=REPO_ID,
        path_in_repo=SUBFOLDER,
        token=HF_TOKEN,
    )
    
    print("✅ Upload complete")

# -----------------------------------------------------------------------------
if __name__ == "__main__":
    print(f"PyTorch: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    export()