lux/modules/capabilities/transcription.nix

{ config, ... }:
let
  homeModules = config.flake.modules.homeManager;
in
{
  flake.modules.homeManager.transcription =
    {
      config,
      lib,
      pkgs,
      ...
    }:
    let
      modelName = "large-v3-turbo-q5_0";
      modelHash = "sha256-OUIhcJzVrR9AxG5gMcphvOiJMebgiMGIKUxtWlX/p+I=";
      modelSha256 = builtins.convertHash {
        hash = modelHash;
        toHashFormat = "base16";
      };
      diarizationModelName = "pyannote/speaker-diarization-community-1";
      diarizationModelRevision = "3533c8cf8e369892e6b79ff1bf80f7b0286a54ee";
      diarizationModel = "${config.xdg.dataHome}/transcribe-recordings/speaker-diarization-community-1";
      model = pkgs.fetchurl {
        name = "ggml-${modelName}.bin";
        url = "https://huggingface.co/ggerganov/whisper.cpp/resolve/98aa99a0a9db05ae2342309f5096248665f7cba3/ggml-${modelName}.bin";
        hash = modelHash;
      };
      diarizationPython = pkgs.python3.withPackages (pythonPackages: [
        pythonPackages.pyannote-audio
      ]);
      diarizeTranscript = pkgs.writeTextFile {
        name = "diarize-transcript";
        destination = "/bin/diarize-transcript";
        executable = true;
        text = ''
          #!${diarizationPython}/bin/python3
          ${builtins.readFile ./_diarize-transcript.py}
        '';
      };
      whisper = pkgs.whisper-cpp;
      prepareDiarizationModel = pkgs.writeShellApplication {
        name = "prepare-transcription-diarization-model";
        runtimeInputs = with pkgs; [
          coreutils
          python3Packages.huggingface-hub
        ];
        text = ''
          model_dir=${lib.escapeShellArg diarizationModel}

          if [[ -f "$model_dir/.complete" ]]; then
            printf 'Speaker diarization model is already installed at %s\n' "$model_dir"
            exit 0
          fi

          printf 'Accept the model terms at https://huggingface.co/${diarizationModelName}\n'
          if ! hf auth whoami >/dev/null 2>&1; then
            hf auth login
          fi

          mkdir -p "$model_dir"
          hf download \
            ${lib.escapeShellArg diarizationModelName} \
            --revision ${lib.escapeShellArg diarizationModelRevision} \
            --local-dir "$model_dir"
          touch "$model_dir/.complete"
          printf 'Installed speaker diarization model at %s\n' "$model_dir"
        '';
      };
      transcribeRecordings = pkgs.writeShellApplication {
        name = "transcribe-recordings";
        runtimeInputs = with pkgs; [
          coreutils
          diarizeTranscript
          ffmpeg
          findutils
          whisper
        ];
        text = builtins.readFile ./_transcribe-recordings.sh;
      };
      home = config.home.homeDirectory;
    in
    {
      home.packages = [
        pkgs.ffmpeg
        prepareDiarizationModel
        whisper
      ];

      home.activation.createTranscriptionLayout = lib.hm.dag.entryAfter [ "writeBoundary" ] ''
        ${pkgs.coreutils}/bin/mkdir -p \
          "${home}/sync/notes/transcripts" \
          "${home}/sync/recordings/inbox" \
          "${home}/sync/recordings/processing" \
          "${home}/sync/recordings/archive" \
          "${home}/sync/recordings/failed"
      '';

      systemd.user.services.transcribe-recordings = {
        Unit.Description = "Transcribe synchronized recordings with whisper.cpp";

        Service = {
          Type = "oneshot";
          ExecStart = lib.getExe transcribeRecordings;
          Environment = [
            "TRANSCRIBE_ENGINE_ACCELERATION=CPU"
            "TRANSCRIBE_ENGINE_VERSION=${whisper.version}"
            "TRANSCRIBE_MODEL=${model}"
            ''"TRANSCRIBE_MODEL_NAME=${modelName} multilingual"''
            "TRANSCRIBE_MODEL_SHA256=${modelSha256}"
            "TRANSCRIBE_DIARIZATION_ENGINE_VERSION=${pkgs.python3Packages.pyannote-audio.version}"
            "TRANSCRIBE_DIARIZATION_MODEL=${diarizationModel}"
            "TRANSCRIBE_DIARIZATION_MODEL_NAME=${diarizationModelName}"
            "TRANSCRIBE_DIARIZATION_MODEL_REVISION=${diarizationModelRevision}"
            "TRANSCRIBE_SETTLE_SECONDS=300"
            "TRANSCRIBE_SYNC_ROOT=${home}/sync"
          ];
          Restart = "on-failure";
          RestartSec = 60;
        };

        Install.WantedBy = [ "default.target" ];
      };

      systemd.user.paths.transcribe-recordings = {
        Unit.Description = "Watch for synchronized recordings";

        Path = {
          PathChanged = "${home}/sync/recordings/inbox";
          Unit = "transcribe-recordings.service";
        };

        Install.WantedBy = [ "paths.target" ];
      };
    };

  flake.modules.nixos.transcription =
    { ... }:
    {
      home-manager.sharedModules = [ homeModules.transcription ];
    };
}