script.py

import os
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"

HF_TOKEN = os.environ.get("HF_TOKEN")
HF_DATASET = os.environ.get("HF_DATASET")
#HF_DATASET = "DevWild/autotrain-pr0b0rk"
repo_id = os.environ.get("MODEL_REPO_ID")

from huggingface_hub import snapshot_download, delete_repo, metadata_update
import uuid
import json
import yaml
import subprocess
import sys
from typing import Optional

#from huggingface_hub import login
#HF_TOKEN = os.getenv("HF_TOKEN")
#if HF_TOKEN:
#    login(token=HF_TOKEN)
#else:
#    raise ValueError("HF_TOKEN environment variable not found!")


if not HF_TOKEN:
    raise ValueError("Missing HF_TOKEN")

if not HF_DATASET:
    raise ValueError("Missing HF_DATASET")

if not repo_id:
    raise ValueError("Missing MODEL_REPO_ID")

# Prevent running script.py twice
LOCKFILE = "/tmp/.script_lock"
if os.path.exists(LOCKFILE):
    print("🔁 Script already ran once — skipping.")
    exit(0)

with open(LOCKFILE, "w") as f:
    f.write("lock")

print("🚀 Running script for the first time")

# START logging
print("🚀 ENV DEBUG START")
print("HF_TOKEN present?", bool(HF_TOKEN))
print("HF_DATASET:", HF_DATASET)
print("MODEL_REPO_ID:", repo_id)
print("🚀 ENV DEBUG END")


#dataset_dir = snapshot_download(HF_DATASET, token=HF_TOKEN)


def download_dataset(hf_dataset_path: str):
    random_id = str(uuid.uuid4())
    snapshot_download(
        repo_id=hf_dataset_path,
        token=HF_TOKEN,
        local_dir=f"/tmp/{random_id}",
        repo_type="dataset",
    )
    return f"/tmp/{random_id}"


def process_dataset(dataset_dir: str):
    # dataset dir consists of images, config.yaml and a metadata.jsonl (optional) with fields: file_name, prompt
    # generate .txt files with the same name as the images with the prompt as the content
    # remove metadata.jsonl
    # return the path to the processed dataset

    # check if config.yaml exists
    if not os.path.exists(os.path.join(dataset_dir, "config.yaml")):
        raise ValueError("config.yaml does not exist")

    # check if metadata.jsonl exists
    if os.path.exists(os.path.join(dataset_dir, "metadata.jsonl")):
        metadata = []
        with open(os.path.join(dataset_dir, "metadata.jsonl"), "r") as f:
            for line in f:
                if len(line.strip()) > 0:
                    metadata.append(json.loads(line))
        for item in metadata:
            txt_path = os.path.join(dataset_dir, item["file_name"])
            txt_path = txt_path.rsplit(".", 1)[0] + ".txt"
            with open(txt_path, "w") as f:
                f.write(item["prompt"])

        # remove metadata.jsonl
        os.remove(os.path.join(dataset_dir, "metadata.jsonl"))

    with open(os.path.join(dataset_dir, "config.yaml"), "r") as f:
        config = yaml.safe_load(f)

    # update config with new dataset
    config["config"]["process"][0]["datasets"][0]["folder_path"] = dataset_dir

    with open(os.path.join(dataset_dir, "config.yaml"), "w") as f:
        yaml.dump(config, f)

    return dataset_dir


def run_training(hf_dataset_path: str):

    dataset_dir = download_dataset(hf_dataset_path)
    dataset_dir = process_dataset(dataset_dir)
    # Force repo_id override in config.yaml
    config_path = os.path.join(dataset_dir, "config.yaml")
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    config["config"]["process"][0]["save"]["hf_repo_id"] = repo_id

    with open(config_path, "w") as f:
        yaml.dump(config, f)

    print("✅ Updated config.yaml with MODEL_REPO_ID:", repo_id)


    # run training
    if not os.path.exists("ai-toolkit"):
        commands = "git clone https://github.com/DevW1ld/ai-toolkit.git ai-toolkit && cd ai-toolkit && git submodule update --init --recursive"
        shutil.rmtree(os.path.join(toolkit_src, ".git"), ignore_errors=True)
        shutil.rmtree(os.path.join(toolkit_src, ".gitmodules"), ignore_errors=True)
        subprocess.run(commands, shell=True)




    # patch_ai_toolkit_typing()
    commands = f"python run.py {os.path.join(dataset_dir, 'config.yaml')}"
    process = subprocess.Popen(commands, shell=True, cwd="ai-toolkit", env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,)

    # Stream logs to Space output
    for line in process.stdout:
        sys.stdout.write(line)
        sys.stdout.flush()
    
    return process, dataset_dir

#def patch_ai_toolkit_typing():
#    config_path = "ai-toolkit/toolkit/config_modules.py"
#    if os.path.exists(config_path):
#        with open(config_path, "r") as f:
#            content = f.read()

#        content = content.replace("torch.Tensor | None", "Optional[torch.Tensor]")

#        with open(config_path, "w") as f:
#            f.write(content)
#        print("✅ Patched ai-toolkit typing for torch.Tensor | None → Optional[torch.Tensor]")
#    else:
#        print("⚠️ Could not patch config_modules.py — file not found")


if __name__ == "__main__":
    try:
        process, dataset_dir = run_training(HF_DATASET)
    #    process.wait()  # Wait for the training process to finish
        exit_code = process.wait()
        print("Training finished with exit code:", exit_code)

        if exit_code != 0:
            raise RuntimeError(f"Training failed with exit code {exit_code}")

        with open(os.path.join(dataset_dir, "config.yaml"), "r") as f:
            config = yaml.safe_load(f)
        #repo_id = config["config"]["process"][0]["save"]["hf_repo_id"]
        #repo_id = os.environ.get("MODEL_REPO_ID")
        #repo_id = os.getenv("MODEL_REPO_ID")
        #repo_id = "DevWild/suppab0rk"

        metadata = {
            "tags": [
                "autotrain",
                "spacerunner",
                "text-to-image",
                "flux",
                "lora",
                "diffusers",
                "template:sd-lora",
            ]
        }    
        metadata_update(repo_id, metadata, token=HF_TOKEN, repo_type="model", overwrite=True)

    finally:
        #delete_repo(HF_DATASET, token=HF_TOKEN, repo_type="dataset", missing_ok=True)
        print("SCRIPT FINISHED, DATASET SHOULD BE DELETED")