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


    

def convert_checkpoint(url, name,repo_id, hf_token ,image_size, scheduler_type, use_half):
    try:
       
        print("Downloading")
        # Download the file
        os.system(f"wget -q {url} --content-disposition -O {name}.safetensors")

        time.sleep(5)
        print("Download successful")
        
        # 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
        
        
        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"{repo_id}/{name}")
        
        # Upload a folder to the repository
        api.upload_folder(
            folder_path=dump_path,
            repo_id=f"{repo_id}/{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="Repo id"),
    #    gr.inputs.Dropdown(label="Visibility", choices=["True","False"]),
        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(),
    title="**Forked from https://huggingface.co./spaces/Androidonnxfork/CivitAi-to-Diffusers**",
    max_queue_size=5
)
iface.launch()