piealamodewhitebread/SillyTavern-Extras1

1import argparse2import base643import gc4import hashlib5import os6import secrets7import sys8import time9import unicodedata10from functools import wraps11from io import BytesIO12from random import randint13 14import markdown15import torch16import webuiapi17from colorama import Fore, Style18from colorama import init as colorama_init19from flask import (Flask, Response, abort, jsonify, render_template_string,20                   request, send_file, send_from_directory)21from flask_compress import Compress22from flask_cors import CORS23from PIL import Image24from transformers import (AutoModelForCausalLM, AutoModelForSeq2SeqLM,25                          AutoProcessor, AutoTokenizer,26                          BlipForConditionalGeneration, pipeline)27 28from constants import *29 30colorama_init()31 32if sys.hexversion < 0x030b0000:33    print(f"{Fore.BLUE}{Style.BRIGHT}Python 3.11 or newer is recommended to run this program.{Style.RESET_ALL}")34    time.sleep(2)35 36class SplitArgs(argparse.Action):37    def __call__(self, parser, namespace, values, option_string=None):38        setattr(39            namespace, self.dest, values.replace('"', "").replace("'", "").split(",")40        )41 42#Setting Root Folders for Silero Generations so it is compatible with STSL, should not effect regular runs. - Rolyat43parent_dir = os.path.dirname(os.path.abspath(__file__))44SILERO_SAMPLES_PATH = os.path.join(parent_dir, "tts_samples")45SILERO_SAMPLE_TEXT = os.path.join(parent_dir)46 47# Create directories if they don't exist48if not os.path.exists(SILERO_SAMPLES_PATH):49    os.makedirs(SILERO_SAMPLES_PATH)50if not os.path.exists(SILERO_SAMPLE_TEXT):51    os.makedirs(SILERO_SAMPLE_TEXT)52 53# Script arguments54parser = argparse.ArgumentParser(55    prog="SillyTavern Extras", description="Web API for transformers models"56)57parser.add_argument(58    "--port", type=int, help="Specify the port on which the application is hosted"59)60parser.add_argument(61    "--listen", action="store_true", help="Host the app on the local network"62)63parser.add_argument(64    "--share", action="store_true", help="Share the app on CloudFlare tunnel"65)66parser.add_argument("--cpu", action="store_true", help="Run the models on the CPU")67parser.add_argument("--cuda", action="store_false", dest="cpu", help="Run the models on the GPU")68parser.add_argument("--cuda-device", help="Specify the CUDA device to use")69parser.add_argument("--mps", "--apple", "--m1", "--m2", action="store_false", dest="cpu", help="Run the models on Apple Silicon")70parser.set_defaults(cpu=True)71parser.add_argument("--summarization-model", help="Load a custom summarization model")72parser.add_argument(73    "--classification-model", help="Load a custom text classification model"74)75parser.add_argument("--captioning-model", help="Load a custom captioning model")76parser.add_argument("--embedding-model", help="Load a custom text embedding model")77parser.add_argument("--chroma-host", help="Host IP for a remote ChromaDB instance")78parser.add_argument("--chroma-port", help="HTTP port for a remote ChromaDB instance (defaults to 8000)")79parser.add_argument("--chroma-folder", help="Path for chromadb persistence folder", default='.chroma_db')80parser.add_argument('--chroma-persist', help="ChromaDB persistence", default=True, action=argparse.BooleanOptionalAction)81parser.add_argument(82    "--secure", action="store_true", help="Enforces the use of an API key"83)84parser.add_argument("--talkinghead-gpu", action="store_true", help="Run the talkinghead animation on the GPU (CPU is default)")85 86parser.add_argument("--coqui-gpu", action="store_true", help="Run the voice models on the GPU (CPU is default)")87parser.add_argument("--coqui-models", help="Install given Coqui-api TTS model at launch (comma separated list, last one will be loaded at start)")88 89parser.add_argument("--max-content-length", help="Set the max")90parser.add_argument("--rvc-save-file", action="store_true", help="Save the last rvc input/output audio file into data/tmp/ folder (for research)")91 92parser.add_argument("--stt-vosk-model-path", help="Load a custom vosk speech-to-text model")93parser.add_argument("--stt-whisper-model-path", help="Load a custom vosk speech-to-text model")94sd_group = parser.add_mutually_exclusive_group()95 96local_sd = parser.add_argument_group("sd-local")97local_sd.add_argument("--sd-model", help="Load a custom SD image generation model")98local_sd.add_argument("--sd-cpu", help="Force the SD pipeline to run on the CPU", action="store_true")99 100remote_sd = parser.add_argument_group("sd-remote")101remote_sd.add_argument(102    "--sd-remote", action="store_true", help="Use a remote backend for SD"103)104remote_sd.add_argument(105    "--sd-remote-host", type=str, help="Specify the host of the remote SD backend"106)107remote_sd.add_argument(108    "--sd-remote-port", type=int, help="Specify the port of the remote SD backend"109)110remote_sd.add_argument(111    "--sd-remote-ssl", action="store_true", help="Use SSL for the remote SD backend"112)113remote_sd.add_argument(114    "--sd-remote-auth",115    type=str,116    help="Specify the username:password for the remote SD backend (if required)",117)118 119parser.add_argument(120    "--enable-modules",121    action=SplitArgs,122    default=[],123    help="Override a list of enabled modules",124)125 126args = parser.parse_args()127# [HF, Huggingface] Set port to 7860, set host to remote.128port = 7860129host = "0.0.0.0"130summarization_model = (131    args.summarization_model132    if args.summarization_model133    else DEFAULT_SUMMARIZATION_MODEL134)135classification_model = (136    args.classification_model137    if args.classification_model138    else DEFAULT_CLASSIFICATION_MODEL139)140captioning_model = (141    args.captioning_model if args.captioning_model else DEFAULT_CAPTIONING_MODEL142)143embedding_model = (144    args.embedding_model if args.embedding_model else DEFAULT_EMBEDDING_MODEL145)146 147sd_use_remote = False if args.sd_model else True148sd_model = args.sd_model if args.sd_model else DEFAULT_SD_MODEL149sd_remote_host = args.sd_remote_host if args.sd_remote_host else DEFAULT_REMOTE_SD_HOST150sd_remote_port = args.sd_remote_port if args.sd_remote_port else DEFAULT_REMOTE_SD_PORT151sd_remote_ssl = args.sd_remote_ssl152sd_remote_auth = args.sd_remote_auth153 154modules = (155    args.enable_modules if args.enable_modules and len(args.enable_modules) > 0 else []156)157 158if len(modules) == 0:159    print(160        f"{Fore.RED}{Style.BRIGHT}You did not select any modules to run! Choose them by adding an --enable-modules option"161    )162    print(f"Example: --enable-modules=caption,summarize{Style.RESET_ALL}")163 164# Models init165cuda_device = DEFAULT_CUDA_DEVICE if not args.cuda_device else args.cuda_device166device_string = cuda_device if torch.cuda.is_available() and not args.cpu else 'mps' if torch.backends.mps.is_available() and not args.cpu else 'cpu'167device = torch.device(device_string)168torch_dtype = torch.float32 if device_string != cuda_device  else torch.float16169 170if not torch.cuda.is_available() and not args.cpu:171    print(f"{Fore.YELLOW}{Style.BRIGHT}torch-cuda is not supported on this device.{Style.RESET_ALL}")172    if not torch.backends.mps.is_available() and not args.cpu:173        print(f"{Fore.YELLOW}{Style.BRIGHT}torch-mps is not supported on this device.{Style.RESET_ALL}")174 175 176print(f"{Fore.GREEN}{Style.BRIGHT}Using torch device: {device_string}{Style.RESET_ALL}")177 178if "talkinghead" in modules:179    import sys180    import threading181    mode = "cuda" if args.talkinghead_gpu else "cpu"182    print("Initializing talkinghead pipeline in " + mode + " mode....")183    talkinghead_path = os.path.abspath(os.path.join(os.getcwd(), "talkinghead"))184    sys.path.append(talkinghead_path) # Add the path to the 'tha3' module to the sys.path list185 186    try:187        import talkinghead.tha3.app.app as talkinghead188        from talkinghead import *189        def launch_talkinghead_gui():190            talkinghead.launch_gui(mode, "separable_float")191        #choices=['standard_float', 'separable_float', 'standard_half', 'separable_half'],192        #choices='The device to use for PyTorch ("cuda" for GPU, "cpu" for CPU).'193        talkinghead_thread = threading.Thread(target=launch_talkinghead_gui)194        talkinghead_thread.daemon = True  # Set the thread as a daemon thread195        talkinghead_thread.start()196 197    except ModuleNotFoundError:198        print("Error: Could not import the 'talkinghead' module.")199 200if "caption" in modules:201    print("Initializing an image captioning model...")202    captioning_processor = AutoProcessor.from_pretrained(captioning_model)203    if "blip" in captioning_model:204        captioning_transformer = BlipForConditionalGeneration.from_pretrained(205            captioning_model, torch_dtype=torch_dtype206        ).to(device)207    else:208        captioning_transformer = AutoModelForCausalLM.from_pretrained(209            captioning_model, torch_dtype=torch_dtype210        ).to(device)211 212if "summarize" in modules:213    print("Initializing a text summarization model...")214    summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model)215    summarization_transformer = AutoModelForSeq2SeqLM.from_pretrained(216        summarization_model, torch_dtype=torch_dtype217    ).to(device)218 219if "sd" in modules and not sd_use_remote:220    from diffusers import (EulerAncestralDiscreteScheduler,221                           StableDiffusionPipeline)222 223    print("Initializing Stable Diffusion pipeline...")224    sd_device_string = cuda_device if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'225    sd_device = torch.device(sd_device_string)226    sd_torch_dtype = torch.float32 if sd_device_string != cuda_device else torch.float16227    sd_pipe = StableDiffusionPipeline.from_pretrained(228        sd_model, custom_pipeline="lpw_stable_diffusion", torch_dtype=sd_torch_dtype229    ).to(sd_device)230    sd_pipe.safety_checker = lambda images, clip_input: (images, False)231    sd_pipe.enable_attention_slicing()232    # pipe.scheduler = KarrasVeScheduler.from_config(pipe.scheduler.config)233    sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(234        sd_pipe.scheduler.config235    )236elif "sd" in modules and sd_use_remote:237    print("Initializing Stable Diffusion connection")238    try:239        sd_remote = webuiapi.WebUIApi(240            host=sd_remote_host, port=sd_remote_port, use_https=sd_remote_ssl241        )242        if sd_remote_auth:243            username, password = sd_remote_auth.split(":")244            sd_remote.set_auth(username, password)245        sd_remote.util_wait_for_ready()246    except Exception as e:247        # remote sd from modules248        print(249            f"{Fore.RED}{Style.BRIGHT}Could not connect to remote SD backend at http{'s' if sd_remote_ssl else ''}://{sd_remote_host}:{sd_remote_port}! Disabling SD module...{Style.RESET_ALL}"250        )251        modules.remove("sd")252 253if "tts" in modules:254    print("tts module is deprecated. Please use silero-tts instead.")255    modules.remove("tts")256    modules.append("silero-tts")257 258 259if "silero-tts" in modules:260    if not os.path.exists(SILERO_SAMPLES_PATH):261        os.makedirs(SILERO_SAMPLES_PATH)262    print("Initializing Silero TTS server")263    from silero_api_server import tts264 265    tts_service = tts.SileroTtsService(SILERO_SAMPLES_PATH)266    if len(os.listdir(SILERO_SAMPLES_PATH)) == 0:267        print("Generating Silero TTS samples...")268        tts_service.update_sample_text(SILERO_SAMPLE_TEXT)269        tts_service.generate_samples()270 271if "edge-tts" in modules:272    print("Initializing Edge TTS client")273    import tts_edge as edge274 275 276if "chromadb" in modules:277    print("Initializing ChromaDB")278    import chromadb279    import posthog280    from chromadb.config import Settings281    from sentence_transformers import SentenceTransformer282 283    # Assume that the user wants in-memory unless a host is specified284    # Also disable chromadb telemetry285    posthog.capture = lambda *args, **kwargs: None286    if args.chroma_host is None:287        if args.chroma_persist:288            chromadb_client = chromadb.PersistentClient(path=args.chroma_folder, settings=Settings(anonymized_telemetry=False))289            print(f"ChromaDB is running in-memory with persistence. Persistence is stored in {args.chroma_folder}. Can be cleared by deleting the folder or purging db.")290        else:291            chromadb_client = chromadb.EphemeralClient(Settings(anonymized_telemetry=False))292            print(f"ChromaDB is running in-memory without persistence.")293    else:294        chroma_port=(295            args.chroma_port if args.chroma_port else DEFAULT_CHROMA_PORT296        )297        chromadb_client = chromadb.HttpClient(host=args.chroma_host, port=chroma_port, settings=Settings(anonymized_telemetry=False))298        print(f"ChromaDB is remotely configured at {args.chroma_host}:{chroma_port}")299 300    chromadb_embedder = SentenceTransformer(embedding_model, device=device_string)301    chromadb_embed_fn = lambda *args, **kwargs: chromadb_embedder.encode(*args, **kwargs).tolist()302 303    # Check if the db is connected and running, otherwise tell the user304    try:305        chromadb_client.heartbeat()306        print("Successfully pinged ChromaDB! Your client is successfully connected.")307    except:308        print("Could not ping ChromaDB! If you are running remotely, please check your host and port!")309 310# Flask init311app = Flask(__name__)312CORS(app)  # allow cross-domain requests313Compress(app) # compress responses314app.config["MAX_CONTENT_LENGTH"] = 500 * 1024 * 1024315 316max_content_length = (317    args.max_content_length318    if args.max_content_length319    else None)320 321if max_content_length is not None:322    print("Setting MAX_CONTENT_LENGTH to",max_content_length,"Mb")323    app.config["MAX_CONTENT_LENGTH"] = int(max_content_length) * 1024 * 1024324 325if "classify" in modules:326    import modules.classify.classify_module as classify_module327    classify_module.init_text_emotion_classifier(classification_model, device, torch_dtype)328 329if "vosk-stt" in modules:330    print("Initializing Vosk speech-recognition (from ST request file)")331    vosk_model_path = (332    args.stt_vosk_model_path333    if args.stt_vosk_model_path334    else None)335 336    import modules.speech_recognition.vosk_module as vosk_module337 338    vosk_module.model = vosk_module.load_model(file_path=vosk_model_path)339    app.add_url_rule("/api/speech-recognition/vosk/process-audio", view_func=vosk_module.process_audio, methods=["POST"])340 341if "whisper-stt" in modules:342    print("Initializing Whisper speech-recognition (from ST request file)")343    whisper_model_path = (344    args.stt_whisper_model_path345    if args.stt_whisper_model_path346    else None)347 348    import modules.speech_recognition.whisper_module as whisper_module349 350    whisper_module.model = whisper_module.load_model(file_path=whisper_model_path)351    app.add_url_rule("/api/speech-recognition/whisper/process-audio", view_func=whisper_module.process_audio, methods=["POST"])352 353if "streaming-stt" in modules:354    print("Initializing vosk/whisper speech-recognition (from extras server microphone)")355    whisper_model_path = (356    args.stt_whisper_model_path357    if args.stt_whisper_model_path358    else None)359 360    import modules.speech_recognition.streaming_module as streaming_module361 362    streaming_module.whisper_model, streaming_module.vosk_model = streaming_module.load_model(file_path=whisper_model_path)363    app.add_url_rule("/api/speech-recognition/streaming/record-and-transcript", view_func=streaming_module.record_and_transcript, methods=["POST"])364 365if "rvc" in modules:366    print("Initializing RVC voice conversion (from ST request file)")367    print("Increasing server upload limit")368    rvc_save_file = (369    args.rvc_save_file370    if args.rvc_save_file371    else False)372 373    if rvc_save_file:374        print("RVC saving file option detected, input/output audio will be savec into data/tmp/ folder")375 376    import sys377    sys.path.insert(0,'modules/voice_conversion')378 379    import modules.voice_conversion.rvc_module as rvc_module380    rvc_module.save_file = rvc_save_file381 382    if "classify" in modules:383        rvc_module.classification_mode = True384 385    rvc_module.fix_model_install()386    app.add_url_rule("/api/voice-conversion/rvc/get-models-list", view_func=rvc_module.rvc_get_models_list, methods=["POST"])387    app.add_url_rule("/api/voice-conversion/rvc/upload-models", view_func=rvc_module.rvc_upload_models, methods=["POST"])388    app.add_url_rule("/api/voice-conversion/rvc/process-audio", view_func=rvc_module.rvc_process_audio, methods=["POST"])389 390 391if "coqui-tts" in modules:392    mode = "GPU" if args.coqui_gpu else "CPU"393    print("Initializing Coqui TTS client in " + mode + " mode")394    import modules.text_to_speech.coqui.coqui_module as coqui_module395 396    if mode == "GPU":397        coqui_module.gpu_mode = True398 399    coqui_models = (400    args.coqui_models401    if args.coqui_models402    else None403    )404 405    if coqui_models is not None:406        coqui_models = coqui_models.split(",")407        for i in coqui_models:408            if not coqui_module.install_model(i):409                raise ValueError("Coqui model loading failed, most likely a wrong model name in --coqui-models argument, check log above to see which one")410 411    # Coqui-api models412    app.add_url_rule("/api/text-to-speech/coqui/coqui-api/check-model-state", view_func=coqui_module.coqui_check_model_state, methods=["POST"])413    app.add_url_rule("/api/text-to-speech/coqui/coqui-api/install-model", view_func=coqui_module.coqui_install_model, methods=["POST"])414 415    # Users models416    app.add_url_rule("/api/text-to-speech/coqui/local/get-models", view_func=coqui_module.coqui_get_local_models, methods=["POST"])417 418    # Handle both coqui-api/users models419    app.add_url_rule("/api/text-to-speech/coqui/generate-tts", view_func=coqui_module.coqui_generate_tts, methods=["POST"])420 421def require_module(name):422    def wrapper(fn):423        @wraps(fn)424        def decorated_view(*args, **kwargs):425            if name not in modules:426                abort(403, "Module is disabled by config")427            return fn(*args, **kwargs)428 429        return decorated_view430 431    return wrapper432 433 434# AI stuff435def classify_text(text: str) -> list:436    return classify_module.classify_text_emotion(text)437 438 439def caption_image(raw_image: Image, max_new_tokens: int = 20) -> str:440    inputs = captioning_processor(raw_image.convert("RGB"), return_tensors="pt").to(441        device, torch_dtype442    )443    outputs = captioning_transformer.generate(**inputs, max_new_tokens=max_new_tokens)444    caption = captioning_processor.decode(outputs[0], skip_special_tokens=True)445    return caption446 447 448def summarize_chunks(text: str, params: dict) -> str:449    try:450        return summarize(text, params)451    except IndexError:452        print(453            "Sequence length too large for model, cutting text in half and calling again"454        )455        new_params = params.copy()456        new_params["max_length"] = new_params["max_length"] // 2457        new_params["min_length"] = new_params["min_length"] // 2458        return summarize_chunks(459            text[: (len(text) // 2)], new_params460        ) + summarize_chunks(text[(len(text) // 2) :], new_params)461 462 463def summarize(text: str, params: dict) -> str:464    # Tokenize input465    inputs = summarization_tokenizer(text, return_tensors="pt").to(device)466    token_count = len(inputs[0])467 468    bad_words_ids = [469        summarization_tokenizer(bad_word, add_special_tokens=False).input_ids470        for bad_word in params["bad_words"]471    ]472    summary_ids = summarization_transformer.generate(473        inputs["input_ids"],474        num_beams=2,475        max_new_tokens=max(token_count, int(params["max_length"])),476        min_new_tokens=min(token_count, int(params["min_length"])),477        repetition_penalty=float(params["repetition_penalty"]),478        temperature=float(params["temperature"]),479        length_penalty=float(params["length_penalty"]),480        bad_words_ids=bad_words_ids,481    )482    summary = summarization_tokenizer.batch_decode(483        summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True484    )[0]485    summary = normalize_string(summary)486    return summary487 488 489def normalize_string(input: str) -> str:490    output = " ".join(unicodedata.normalize("NFKC", input).strip().split())491    return output492 493 494def generate_image(data: dict) -> Image:495    prompt = normalize_string(f'{data["prompt_prefix"]} {data["prompt"]}')496 497    if sd_use_remote:498        image = sd_remote.txt2img(499            prompt=prompt,500            negative_prompt=data["negative_prompt"],501            sampler_name=data["sampler"],502            steps=data["steps"],503            cfg_scale=data["scale"],504            width=data["width"],505            height=data["height"],506            restore_faces=data["restore_faces"],507            enable_hr=data["enable_hr"],508            save_images=True,509            send_images=True,510            do_not_save_grid=False,511            do_not_save_samples=False,512        ).image513    else:514        image = sd_pipe(515            prompt=prompt,516            negative_prompt=data["negative_prompt"],517            num_inference_steps=data["steps"],518            guidance_scale=data["scale"],519            width=data["width"],520            height=data["height"],521        ).images[0]522 523    image.save("./debug.png")524    return image525 526 527def image_to_base64(image: Image, quality: int = 75) -> str:528    buffer = BytesIO()529    image.convert("RGB")530    image.save(buffer, format="JPEG", quality=quality)531    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")532    return img_str533 534 535ignore_auth = []536# [HF, Huggingface] Get password instead of text file.537api_key = os.environ.get("password")538 539def is_authorize_ignored(request):540    view_func = app.view_functions.get(request.endpoint)541 542    if view_func is not None:543        if view_func in ignore_auth:544            return True545    return False546 547 548@app.before_request549def before_request():550    # Request time measuring551    request.start_time = time.time()552 553    # Checks if an API key is present and valid, otherwise return unauthorized554    # The options check is required so CORS doesn't get angry555    try:556        if request.method != 'OPTIONS' and is_authorize_ignored(request) == False and getattr(request.authorization, 'token', '') != api_key:557            print(f"WARNING: Unauthorized API key access from {request.remote_addr}")558            if request.method == 'POST':559                print(f"Incoming POST request with {request.headers.get('Authorization')}")560            response = jsonify({ 'error': '401: Invalid API key' })561            response.status_code = 401562            return "https://(hf_name)-(space_name).hf.space/"563    except Exception as e:564        print(f"API key check error: {e}")565        return "https://(hf_name)-(space_name).hf.space/"566 567 568@app.after_request569def after_request(response):570    duration = time.time() - request.start_time571    response.headers["X-Request-Duration"] = str(duration)572    return response573 574 575@app.route("/", methods=["GET"])576def index():577    with open("./README.md", "r", encoding="utf8") as f:578        content = f.read()579    return render_template_string(markdown.markdown(content, extensions=["tables"]))580 581 582@app.route("/api/extensions", methods=["GET"])583def get_extensions():584    extensions = dict(585        {586            "extensions": [587                {588                    "name": "not-supported",589                    "metadata": {590                        "display_name": """<span style="white-space:break-spaces;">Extensions serving using Extensions API is no longer supported. Please update the mod from: <a href="https://github.com/Cohee1207/SillyTavern">https://github.com/Cohee1207/SillyTavern</a></span>""",591                        "requires": [],592                        "assets": [],593                    },594                }595            ]596        }597    )598    return jsonify(extensions)599 600 601@app.route("/api/caption", methods=["POST"])602@require_module("caption")603def api_caption():604    data = request.get_json()605 606    if "image" not in data or not isinstance(data["image"], str):607        abort(400, '"image" is required')608 609    image = Image.open(BytesIO(base64.b64decode(data["image"])))610    image = image.convert("RGB")611    image.thumbnail((512, 512))612    caption = caption_image(image)613    thumbnail = image_to_base64(image)614    print("Caption:", caption, sep="\n")615    gc.collect()616    return jsonify({"caption": caption, "thumbnail": thumbnail})617 618 619@app.route("/api/summarize", methods=["POST"])620@require_module("summarize")621def api_summarize():622    data = request.get_json()623 624    if "text" not in data or not isinstance(data["text"], str):625        abort(400, '"text" is required')626 627    params = DEFAULT_SUMMARIZE_PARAMS.copy()628 629    if "params" in data and isinstance(data["params"], dict):630        params.update(data["params"])631 632    print("Summary input:", data["text"], sep="\n")633    summary = summarize_chunks(data["text"], params)634    print("Summary output:", summary, sep="\n")635    gc.collect()636    return jsonify({"summary": summary})637 638 639@app.route("/api/classify", methods=["POST"])640@require_module("classify")641def api_classify():642    data = request.get_json()643 644    if "text" not in data or not isinstance(data["text"], str):645        abort(400, '"text" is required')646 647    print("Classification input:", data["text"], sep="\n")648    classification = classify_text(data["text"])649    print("Classification output:", classification, sep="\n")650    gc.collect()651    if "talkinghead" in modules: #send emotion to talkinghead652        talkinghead.setEmotion(classification)653    return jsonify({"classification": classification})654 655 656@app.route("/api/classify/labels", methods=["GET"])657@require_module("classify")658def api_classify_labels():659    classification = classify_text("")660    labels = [x["label"] for x in classification]661    if "talkinghead" in modules:662        labels.append('talkinghead')  # Add 'talkinghead' to the labels list663    return jsonify({"labels": labels})664 665@app.route("/api/talkinghead/load", methods=["POST"])666def live_load():667    file = request.files['file']668    # convert stream to bytes and pass to talkinghead_load669    return talkinghead.talkinghead_load_file(file.stream)670 671@app.route('/api/talkinghead/unload')672def live_unload():673    return talkinghead.unload()674 675@app.route('/api/talkinghead/start_talking')676def start_talking():677    return talkinghead.start_talking()678 679@app.route('/api/talkinghead/stop_talking')680def stop_talking():681    return talkinghead.stop_talking()682 683@app.route('/api/talkinghead/result_feed')684def result_feed():685    return talkinghead.result_feed()686 687@app.route("/api/image", methods=["POST"])688@require_module("sd")689def api_image():690    required_fields = {691        "prompt": str,692    }693 694    optional_fields = {695        "steps": 30,696        "scale": 6,697        "sampler": "DDIM",698        "width": 512,699        "height": 512,700        "restore_faces": False,701        "enable_hr": False,702        "prompt_prefix": PROMPT_PREFIX,703        "negative_prompt": NEGATIVE_PROMPT,704    }705 706    data = request.get_json()707 708    # Check required fields709    for field, field_type in required_fields.items():710        if field not in data or not isinstance(data[field], field_type):711            abort(400, f'"{field}" is required')712 713    # Set optional fields to default values if not provided714    for field, default_value in optional_fields.items():715        type_match = (716            (int, float)717            if isinstance(default_value, (int, float))718            else type(default_value)719        )720        if field not in data or not isinstance(data[field], type_match):721            data[field] = default_value722 723    try:724        print("SD inputs:", data, sep="\n")725        image = generate_image(data)726        base64image = image_to_base64(image, quality=90)727        return jsonify({"image": base64image})728    except RuntimeError as e:729        abort(400, str(e))730 731 732@app.route("/api/image/model", methods=["POST"])733@require_module("sd")734def api_image_model_set():735    data = request.get_json()736 737    if not sd_use_remote:738        abort(400, "Changing model for local sd is not supported.")739    if "model" not in data or not isinstance(data["model"], str):740        abort(400, '"model" is required')741 742    old_model = sd_remote.util_get_current_model()743    sd_remote.util_set_model(data["model"], find_closest=False)744    # sd_remote.util_set_model(data['model'])745    sd_remote.util_wait_for_ready()746    new_model = sd_remote.util_get_current_model()747 748    return jsonify({"previous_model": old_model, "current_model": new_model})749 750 751@app.route("/api/image/model", methods=["GET"])752@require_module("sd")753def api_image_model_get():754    model = sd_model755 756    if sd_use_remote:757        model = sd_remote.util_get_current_model()758 759    return jsonify({"model": model})760 761 762@app.route("/api/image/models", methods=["GET"])763@require_module("sd")764def api_image_models():765    models = [sd_model]766 767    if sd_use_remote:768        models = sd_remote.util_get_model_names()769 770    return jsonify({"models": models})771 772 773@app.route("/api/image/samplers", methods=["GET"])774@require_module("sd")775def api_image_samplers():776    samplers = ["Euler a"]777 778    if sd_use_remote:779        samplers = [sampler["name"] for sampler in sd_remote.get_samplers()]780 781    return jsonify({"samplers": samplers})782 783 784@app.route("/api/modules", methods=["GET"])785def get_modules():786    return jsonify({"modules": modules})787 788 789@app.route("/api/tts/speakers", methods=["GET"])790@require_module("silero-tts")791def tts_speakers():792    voices = [793        {794            "name": speaker,795            "voice_id": speaker,796            "preview_url": f"{str(request.url_root)}api/tts/sample/{speaker}",797        }798        for speaker in tts_service.get_speakers()799    ]800    return jsonify(voices)801 802# Added fix for Silero not working as new files were unable to be created if one already existed. - Rolyat 7/7/23803@app.route("/api/tts/generate", methods=["POST"])804@require_module("silero-tts")805def tts_generate():806    voice = request.get_json()807    if "text" not in voice or not isinstance(voice["text"], str):808        abort(400, '"text" is required')809    if "speaker" not in voice or not isinstance(voice["speaker"], str):810        abort(400, '"speaker" is required')811    # Remove asterisks812    voice["text"] = voice["text"].replace("*", "")813    try:814        # Remove the destination file if it already exists815        if os.path.exists('test.wav'):816            os.remove('test.wav')817 818        audio = tts_service.generate(voice["speaker"], voice["text"])819        audio_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.path.basename(audio))820 821        os.rename(audio, audio_file_path)822        return send_file(audio_file_path, mimetype="audio/x-wav")823    except Exception as e:824        print(e)825        abort(500, voice["speaker"])826 827 828@app.route("/api/tts/sample/<speaker>", methods=["GET"])829@require_module("silero-tts")830def tts_play_sample(speaker: str):831    return send_from_directory(SILERO_SAMPLES_PATH, f"{speaker}.wav")832 833 834@app.route("/api/edge-tts/list", methods=["GET"])835@require_module("edge-tts")836def edge_tts_list():837    voices = edge.get_voices()838    return jsonify(voices)839 840 841@app.route("/api/edge-tts/generate", methods=["POST"])842@require_module("edge-tts")843def edge_tts_generate():844    data = request.get_json()845    if "text" not in data or not isinstance(data["text"], str):846        abort(400, '"text" is required')847    if "voice" not in data or not isinstance(data["voice"], str):848        abort(400, '"voice" is required')849    if "rate" in data and isinstance(data['rate'], int):850        rate = data['rate']851    else:852        rate = 0853    # Remove asterisks854    data["text"] = data["text"].replace("*", "")855    try:856        audio = edge.generate_audio(text=data["text"], voice=data["voice"], rate=rate)857        return Response(audio, mimetype="audio/mpeg")858    except Exception as e:859        print(e)860        abort(500, data["voice"])861 862 863@app.route("/api/chromadb", methods=["POST"])864@require_module("chromadb")865def chromadb_add_messages():866    data = request.get_json()867    if "chat_id" not in data or not isinstance(data["chat_id"], str):868        abort(400, '"chat_id" is required')869    if "messages" not in data or not isinstance(data["messages"], list):870        abort(400, '"messages" is required')871 872    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()873    collection = chromadb_client.get_or_create_collection(874        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn875    )876 877    documents = [m["content"] for m in data["messages"]]878    ids = [m["id"] for m in data["messages"]]879    metadatas = [880        {"role": m["role"], "date": m["date"], "meta": m.get("meta", "")}881        for m in data["messages"]882    ]883 884    collection.upsert(885        ids=ids,886        documents=documents,887        metadatas=metadatas,888    )889 890    return jsonify({"count": len(ids)})891 892 893@app.route("/api/chromadb/purge", methods=["POST"])894@require_module("chromadb")895def chromadb_purge():896    data = request.get_json()897    if "chat_id" not in data or not isinstance(data["chat_id"], str):898        abort(400, '"chat_id" is required')899 900    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()901    collection = chromadb_client.get_or_create_collection(902        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn903    )904 905    count = collection.count()906    collection.delete()907    print("ChromaDB embeddings deleted", count)908    return 'Ok', 200909 910 911@app.route("/api/chromadb/query", methods=["POST"])912@require_module("chromadb")913def chromadb_query():914    data = request.get_json()915    if "chat_id" not in data or not isinstance(data["chat_id"], str):916        abort(400, '"chat_id" is required')917    if "query" not in data or not isinstance(data["query"], str):918        abort(400, '"query" is required')919 920    if "n_results" not in data or not isinstance(data["n_results"], int):921        n_results = 1922    else:923        n_results = data["n_results"]924 925    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()926    collection = chromadb_client.get_or_create_collection(927        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn928    )929 930    if collection.count() == 0:931        print(f"Queried empty/missing collection for {repr(data['chat_id'])}.")932        return jsonify([])933 934 935    n_results = min(collection.count(), n_results)936    query_result = collection.query(937        query_texts=[data["query"]],938        n_results=n_results,939    )940 941    documents = query_result["documents"][0]942    ids = query_result["ids"][0]943    metadatas = query_result["metadatas"][0]944    distances = query_result["distances"][0]945 946    messages = [947        {948            "id": ids[i],949            "date": metadatas[i]["date"],950            "role": metadatas[i]["role"],951            "meta": metadatas[i]["meta"],952            "content": documents[i],953            "distance": distances[i],954        }955        for i in range(len(ids))956    ]957 958    return jsonify(messages)959 960@app.route("/api/chromadb/multiquery", methods=["POST"])961@require_module("chromadb")962def chromadb_multiquery():963    data = request.get_json()964    if "chat_list" not in data or not isinstance(data["chat_list"], list):965        abort(400, '"chat_list" is required and should be a list')966    if "query" not in data or not isinstance(data["query"], str):967        abort(400, '"query" is required')968 969    if "n_results" not in data or not isinstance(data["n_results"], int):970        n_results = 1971    else:972        n_results = data["n_results"]973 974    messages = []975 976    for chat_id in data["chat_list"]:977        if not isinstance(chat_id, str):978            continue979 980        try:981            chat_id_md5 = hashlib.md5(chat_id.encode()).hexdigest()982            collection = chromadb_client.get_collection(983                name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn984            )985 986            # Skip this chat if the collection is empty987            if collection.count() == 0:988                continue989 990            n_results_per_chat = min(collection.count(), n_results)991            query_result = collection.query(992                query_texts=[data["query"]],993                n_results=n_results_per_chat,994            )995            documents = query_result["documents"][0]996            ids = query_result["ids"][0]997            metadatas = query_result["metadatas"][0]998            distances = query_result["distances"][0]999 1000            chat_messages = [1001                {1002                    "id": ids[i],1003                    "date": metadatas[i]["date"],1004                    "role": metadatas[i]["role"],1005                    "meta": metadatas[i]["meta"],1006                    "content": documents[i],1007                    "distance": distances[i],1008                }1009                for i in range(len(ids))1010            ]1011 1012            messages.extend(chat_messages)1013        except Exception as e:1014            print(e)1015 1016    #remove duplicate msgs, filter down to the right number1017    seen = set()1018    messages = [d for d in messages if not (d['content'] in seen or seen.add(d['content']))]1019    messages = sorted(messages, key=lambda x: x['distance'])[0:n_results]1020 1021    return jsonify(messages)1022 1023 1024@app.route("/api/chromadb/export", methods=["POST"])1025@require_module("chromadb")1026def chromadb_export():1027    data = request.get_json()1028    if "chat_id" not in data or not isinstance(data["chat_id"], str):1029        abort(400, '"chat_id" is required')1030 1031    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()1032    try:1033        collection = chromadb_client.get_collection(1034            name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn1035        )1036    except Exception as e:1037        print(e)1038        abort(400, "Chat collection not found in chromadb")1039 1040    collection_content = collection.get()1041    documents = collection_content.get('documents', [])1042    ids = collection_content.get('ids', [])1043    metadatas = collection_content.get('metadatas', [])1044 1045    unsorted_content = [1046        {1047            "id": ids[i],1048            "metadata": metadatas[i],1049            "document": documents[i],1050        }1051        for i in range(len(ids))1052    ]1053 1054    sorted_content = sorted(unsorted_content, key=lambda x: x['metadata']['date'])1055 1056    export = {1057        "chat_id": data["chat_id"],1058        "content": sorted_content1059    }1060 1061    return jsonify(export)1062 1063@app.route("/api/chromadb/import", methods=["POST"])1064@require_module("chromadb")1065def chromadb_import():1066    data = request.get_json()1067    content = data['content']1068    if "chat_id" not in data or not isinstance(data["chat_id"], str):1069        abort(400, '"chat_id" is required')1070 1071    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()1072    collection = chromadb_client.get_or_create_collection(1073        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn1074    )1075 1076    documents = [item['document'] for item in content]1077    metadatas = [item['metadata'] for item in content]1078    ids = [item['id'] for item in content]1079 1080 1081    collection.upsert(documents=documents, metadatas=metadatas, ids=ids)1082    print(f"Imported {len(ids)} (total {collection.count()}) content entries into {repr(data['chat_id'])}")1083 1084    return jsonify({"count": len(ids)})1085 1086 1087if args.share:1088    import inspect1089 1090    from flask_cloudflared import _run_cloudflared1091 1092    sig = inspect.signature(_run_cloudflared)1093    sum = sum(1094        11095        for param in sig.parameters.values()1096        if param.kind == param.POSITIONAL_OR_KEYWORD1097    )1098    if sum > 1:1099        metrics_port = randint(8100, 9000)1100        cloudflare = _run_cloudflared(port, metrics_port)1101    else:1102        cloudflare = _run_cloudflared(port)1103    print(f"{Fore.GREEN}{Style.NORMAL}Running on: {cloudflare}{Style.RESET_ALL}")1104 1105ignore_auth.append(tts_play_sample)1106ignore_auth.append(result_feed)1107app.run(host=host, port=port)1108