UploadModel.tsx

import {
  Alert,
  Button,
  Card,
  CardBanner,
  CardBody,
  CardHeader,
  HStack,
  VStack,
} from "@client/component-library";
import type {
  UploadModelError,
  UploadModelResponse,
} from "@shared/api-types/models";
import { uploadModelErrorSchema } from "@shared/api-types/models";
import { paths } from "@shared/utils/paths";
import { isAxiosError } from "axios";
import { useState } from "react";
import { InputFormField } from "src/components/common/FormFields/InputFormField";
import { TextAreaFormField } from "src/components/common/FormFields/TextAreaFormField";
import { CUSTOM_MODELS_LINK } from "src/constants";
import { getUploadModelCodeSnippets } from "src/utils/codeViewerSnippets";

import { APIReferenceSection } from "./APIReferenceSection";
import { UploadJobProgress } from "./UploadJobProgress";
import { useUploadModel } from "./useUploadModel";

type Model = {
  source: string;
  name: string;
  description: string;
  huggingfaceToken: string;
  hasChangedNameManually: boolean;
};

type IDLE = {
  current: "IDLE";
  model: Model;
  upload(): void;
  updateModel(model: Partial<Model>): void;
};

type UPLOADING = {
  current: "UPLOADING";
  model: Model;
};

type UPLOADED = {
  current: "UPLOADED";
  model: Model;
  uploadJobId: string;
  response: UploadModelResponse;
  completed(newModelName: string): void;
};

type COMPLETED = {
  current: "COMPLETED";
  model: Model;
  newModelName: string;
};

type ERROR = {
  current: "ERROR";
  model: Model;
  type: string;
  message: string;
};

type State = IDLE | UPLOADING | UPLOADED | COMPLETED | ERROR;

export const UploadModelPage = () => {
  const [state, setState] = useState<State>(
    IDLE()
  );

  const { mutate: uploadModel, isPending: isRequestingUpload } =
    useUploadModel();
  const hasMissingFields =
    !state.model.name || !state.model.source || !state.model.huggingfaceToken;
  const showOpenModelButton = state.current === "UPLOADED";
  const resultModelName =
    state.current === "UPLOADED"
      ? state.response.model_name
      : state.model.name;

  return <div />;

  function IDLE(model: Model = {
    source: "",
    name: "",
    description: "",
    huggingfaceToken: "",
    hasChangedNameManually: false,
  }): IDLE {
    return {
      current: "IDLE",
      model,
      upload() {
        if (!model.name || !model.source || !model.huggingfaceToken) {
          return;
        }

        setState(UPLOADING(model));
      },
      updateModel(updatedModel) {
        setState(IDLE(updateModel(model, updatedModel)));
      },
    };
  }

  function UPLOADING(model: Model): UPLOADING {
    uploadModel(
      {
        model_name: model.name,
        model_source: model.source,
        description: model.description,
        hf_token: model.huggingfaceToken,
      },
      {
        onSuccess(response) {
          const uploadJobId = response.data?.job_id ?? null;

          setState(uploadJobId ? UPLOADED(model, uploadJobId, response) : IDLE());
        },
        onError(error) {
          const customError =
            isAxiosError(error) &&
            uploadModelErrorSchema.safeParse(error.response?.data);

          setState(
            customError.success
              ? ERROR(model, customError.data.type, customError.data.message)
              : ERROR(
                  model,
                  "unknown_error",
                  "Unexpected error. Please contact support if the problem persists."
                )
          );
        },
      }
    );

    return {
      current: "UPLOADING",
      model,
    };
  }

  function UPLOADED(
    model: Model,
    uploadJobId: string,
    response: UploadModelResponse
  ): UPLOADED {
    return {
      current: "UPLOADED",
      model,
      uploadJobId,
      response,
      completed(newModelName) {
        setState(COMPLETED(model, newModelName));
      }
    };
  }

  function ERROR(model: Model, type: string, message: string): ERROR {
    return {
      current: "ERROR",
      model,
      type,
      message,
    };
  }
  
  function COMPLETED(model: Model, newModelName: string): COMPLETED {
    return {
      current: "COMPLETED",
      model,
      newModelName,
    };
  }
};

function updateModel(model: Model, updatedModel: Partial<Model>): Model {
  let name = model.name;

  if (updatedModel.name && !model.hasChangedNameManually) {
    // If the user did not touch the name, we automatically generate it
    let orgAndName = name;
    if (name.startsWith("https://huggingface.co/")) {
      orgAndName = name.replace("https://huggingface.co/", "");
    }

    name = orgAndName.split("/").join("-");
  }

  return {
    ...model,
    ...updatedModel,
    name,
  };
}