import gradio as gr
import time
import os
from huggingface_hub import HfApi, create_repo

def convert_checkpoint(url, name, hf_token, image_size, scheduler_type, use_half):
    try:
        # Download the file
        os.system(f"wget {url} --content-disposition -O {name}.safetensors")
        
        # Introduce a delay of 30 seconds
        time.sleep(30)
        
        # Construct the checkpoint path and dump path
        checkpoint_path = f"{name}.safetensors"
        dump_path = f"/home/user/app/{name}"
        
        cmd = [
            "python3",
            "diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py",  # Replace with the name of your script
            "--checkpoint_path", checkpoint_path,
            f"--scheduler_type {scheduler_type}",
            f"--image_size {image_size}",
            "--prediction_type epsilon",
            "--device cpu",
            "--from_safetensors",
            "--to_safetensors",
            "--dump_path", dump_path
        ]
        
        if use_half:
            cmd.append("--half")
        
        result = os.system(" ".join(cmd))
        output = result
        
        # Clean up downloaded file
        os.remove(checkpoint_path)
        
        # Log in to your Hugging Face account
        os.system(f"huggingface-cli login --token {hf_token}")
        
        # Create a repository
        api = HfApi()
        api.create_repo(f"Androidonnxfork/{name}", private=True)
        
        # Upload a folder to the repository
        api.upload_folder(
            folder_path=dump_path,
            repo_id=f"Androidonnxfork/{name}",
            repo_type="model",
        )
        
    except Exception as e:
        output = str(e)
    
    return output

iface = gr.Interface(
    fn=convert_checkpoint,
    inputs=[
        gr.inputs.Textbox(label="URL"),
        gr.inputs.Textbox(label="Name"),
        gr.inputs.Textbox(label="Hugging Face API Token"),
        gr.inputs.Radio(label="Image Size", choices=["512", "768"]),
        gr.inputs.Dropdown(label="Scheduler Type", choices=['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']),
        gr.inputs.Checkbox(label="Use Half Precision")
    ],
    outputs=gr.outputs.Textbox()
)
iface.launch()