Upload 4 files

constants.py

@@ -1,18 +1,18 @@
 # Constants
-# Also try: 'slauw87/bart-large-cnn-samsum'
-DEFAULT_SUMMARIZATION_MODEL = "Qiliang/bart-large-cnn-samsum-ElectrifAi_v14"
-# Also try: 'nateraw/bert-base-uncased-emotion'
-DEFAULT_CLASSIFICATION_MODEL = "joeddav/distilbert-base-uncased-go-emotions-student"
+DEFAULT_CUDA_DEVICE = "cuda:0"
+# Also try: 'Qiliang/bart-large-cnn-samsum-ElectrifAi_v10'
+DEFAULT_SUMMARIZATION_MODEL = "Qiliang/bart-large-cnn-samsum-ChatGPT_v3"
+# Also try: 'joeddav/distilbert-base-uncased-go-emotions-student'
+DEFAULT_CLASSIFICATION_MODEL = "nateraw/bert-base-uncased-emotion"
 # Also try: 'Salesforce/blip-image-captioning-base'
 DEFAULT_CAPTIONING_MODEL = "Salesforce/blip-image-captioning-large"
-# Also try: 'ckpt/anything-v4.5-vae-swapped'
-DEFAULT_SD_MODEL = "sinkinai/MeinaHentai-v3-baked-vae"
+DEFAULT_SD_MODEL = "ckpt/anything-v4.5-vae-swapped"
 DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
 DEFAULT_REMOTE_SD_HOST = "127.0.0.1"
 DEFAULT_REMOTE_SD_PORT = 7860
 DEFAULT_CHROMA_PORT = 8000
 SILERO_SAMPLES_PATH = "tts_samples"
-SILERO_SAMPLE_TEXT = "Doctor is your lord and savior"
+SILERO_SAMPLE_TEXT = "The quick brown fox jumps over the lazy dog"
 # ALL_MODULES = ['caption', 'summarize', 'classify', 'keywords', 'prompt', 'sd']
 DEFAULT_SUMMARIZE_PARAMS = {
     "temperature": 1.0,

requirements-complete.txt

@@ -0,0 +1,19 @@
+flask
+flask-cloudflared
+flask-cors
+flask-compress
+markdown
+Pillow
+colorama
+webuiapi
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==2.0.0+cu117
+torchvision==0.15.1
+torchaudio==2.0.1+cu117
+accelerate
+transformers==4.28.1
+diffusers==0.16.1
+silero-api-server
+chromadb
+sentence_transformers
+edge-tts

server.py

@@ -21,6 +21,7 @@ import torch
 import time
 import os
 import gc
+import sys
 import secrets
 from PIL import Image
 import base64
@@ -33,6 +34,9 @@ from colorama import Fore, Style, init as colorama_init
 
 colorama_init()
 
+if sys.hexversion < 0x030b0000:
+    print(f"{Fore.BLUE}{Style.BRIGHT}Python 3.11 or newer is recommended to run this program.{Style.RESET_ALL}")
+    time.sleep(2)
 
 class SplitArgs(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
@@ -40,6 +44,16 @@ class SplitArgs(argparse.Action):
             namespace, self.dest, values.replace('"', "").replace("'", "").split(",")
         )
 
+#Setting Root Folders for Silero Generations so it is compatible with STSL, should not effect regular runs. - Rolyat
+parent_dir = os.path.dirname(os.path.abspath(__file__))
+SILERO_SAMPLES_PATH = os.path.join(parent_dir, "tts_samples")
+SILERO_SAMPLE_TEXT = os.path.join(parent_dir)
+
+# Create directories if they don't exist
+if not os.path.exists(SILERO_SAMPLES_PATH):
+    os.makedirs(SILERO_SAMPLES_PATH)
+if not os.path.exists(SILERO_SAMPLE_TEXT):
+    os.makedirs(SILERO_SAMPLE_TEXT)
 
 # Script arguments
 parser = argparse.ArgumentParser(
@@ -56,6 +70,8 @@ parser.add_argument(
 )
 parser.add_argument("--cpu", action="store_true", help="Run the models on the CPU")
 parser.add_argument("--cuda", action="store_false", dest="cpu", help="Run the models on the GPU")
+parser.add_argument("--cuda-device", help="Specify the CUDA device to use")
+parser.add_argument("--mps", "--apple", "--m1", "--m2", action="store_false", dest="cpu", help="Run the models on Apple Silicon")
 parser.set_defaults(cpu=True)
 parser.add_argument("--summarization-model", help="Load a custom summarization model")
 parser.add_argument(
@@ -66,11 +82,10 @@ parser.add_argument("--embedding-model", help="Load a custom text embedding mode
 parser.add_argument("--chroma-host", help="Host IP for a remote ChromaDB instance")
 parser.add_argument("--chroma-port", help="HTTP port for a remote ChromaDB instance (defaults to 8000)")
 parser.add_argument("--chroma-folder", help="Path for chromadb persistence folder", default='.chroma_db')
-parser.add_argument('--chroma-persist', help="Chromadb persistence", default=True, action=argparse.BooleanOptionalAction)
+parser.add_argument('--chroma-persist', help="ChromaDB persistence", default=True, action=argparse.BooleanOptionalAction)
 parser.add_argument(
     "--secure", action="store_true", help="Enforces the use of an API key"
 )
-
 sd_group = parser.add_mutually_exclusive_group()
 
 local_sd = sd_group.add_argument_group("sd-local")
@@ -105,8 +120,8 @@ parser.add_argument(
 
 args = parser.parse_args()
 
-port = 7860
-host = "0.0.0.0"
+port = args.port if args.port else 5100
+host = "0.0.0.0" if args.listen else "localhost"
 summarization_model = (
     args.summarization_model
     if args.summarization_model
@@ -142,12 +157,16 @@ if len(modules) == 0:
     print(f"Example: --enable-modules=caption,summarize{Style.RESET_ALL}")
 
 # Models init
-device_string = "cuda:0" if torch.cuda.is_available() and not args.cpu else "cpu"
+cuda_device = DEFAULT_CUDA_DEVICE if not args.cuda_device else args.cuda_device
+device_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'
 device = torch.device(device_string)
-torch_dtype = torch.float32 if device_string == "cpu" else torch.float16
+torch_dtype = torch.float32 if device_string != cuda_device  else torch.float16
 
 if not torch.cuda.is_available() and not args.cpu:
-    print(f"{Fore.YELLOW}{Style.BRIGHT}torch-cuda is not supported on this device. Defaulting to CPU mode.{Style.RESET_ALL}")
+    print(f"{Fore.YELLOW}{Style.BRIGHT}torch-cuda is not supported on this device.{Style.RESET_ALL}")
+    if not torch.backends.mps.is_available() and not args.cpu:
+        print(f"{Fore.YELLOW}{Style.BRIGHT}torch-mps is not supported on this device.{Style.RESET_ALL}")
+
 
 print(f"{Fore.GREEN}{Style.BRIGHT}Using torch device: {device_string}{Style.RESET_ALL}")
 
@@ -184,12 +203,10 @@ if "sd" in modules and not sd_use_remote:
     from diffusers import StableDiffusionPipeline
     from diffusers import EulerAncestralDiscreteScheduler
 
-    print("Initializing Stable Diffusion pipeline")
-    sd_device_string = (
-        "cuda" if torch.cuda.is_available() and not args.sd_cpu else "cpu"
-    )
+    print("Initializing Stable Diffusion pipeline...")
+    sd_device_string = cuda_device if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
     sd_device = torch.device(sd_device_string)
-    sd_torch_dtype = torch.float32 if sd_device_string == "cpu" else torch.float16
+    sd_torch_dtype = torch.float32 if sd_device_string != cuda_device else torch.float16
     sd_pipe = StableDiffusionPipeline.from_pretrained(
         sd_model, custom_pipeline="lpw_stable_diffusion", torch_dtype=sd_torch_dtype
     ).to(sd_device)
@@ -252,26 +269,19 @@ if "chromadb" in modules:
     posthog.capture = lambda *args, **kwargs: None
     if args.chroma_host is None:
         if args.chroma_persist:
-            chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False, persist_directory=args.chroma_folder, chroma_db_impl='duckdb+parquet'))
+            chromadb_client = chromadb.PersistentClient(path=args.chroma_folder, settings=Settings(anonymized_telemetry=False))
             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.")
         else:
-            chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
+            chromadb_client = chromadb.EphemeralClient(Settings(anonymized_telemetry=False))
             print(f"ChromaDB is running in-memory without persistence.")
     else:
         chroma_port=(
             args.chroma_port if args.chroma_port else DEFAULT_CHROMA_PORT
         )
-        chromadb_client = chromadb.Client(
-            Settings(
-                anonymized_telemetry=False,
-                chroma_api_impl="rest",
-                chroma_server_host=args.chroma_host,
-                chroma_server_http_port=chroma_port
-            )
-        )
+        chromadb_client = chromadb.HttpClient(host=args.chroma_host, port=chroma_port, settings=Settings(anonymized_telemetry=False))
         print(f"ChromaDB is remotely configured at {args.chroma_host}:{chroma_port}")
 
-    chromadb_embedder = SentenceTransformer(embedding_model)
+    chromadb_embedder = SentenceTransformer(embedding_model, device=device_string)
     chromadb_embed_fn = lambda *args, **kwargs: chromadb_embedder.encode(*args, **kwargs).tolist()
 
     # Check if the db is connected and running, otherwise tell the user
@@ -405,10 +415,24 @@ def image_to_base64(image: Image, quality: int = 75) -> str:
     image.save(buffer, format="JPEG", quality=quality)
     img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
     return img_str
-    
-ignore_auth = []    
 
-api_key = os.environ.get("password")
+ignore_auth = []
+# Reads an API key from an already existing file. If that file doesn't exist, create it.
+if args.secure:
+    try:
+        with open("api_key.txt", "r") as txt:
+            api_key = txt.read().replace('\n', '')
+    except:
+        api_key = secrets.token_hex(5)
+        with open("api_key.txt", "w") as txt:
+            txt.write(api_key)
+
+    print(f"Your API key is {api_key}")
+elif args.share and args.secure != True:
+    print("WARNING: This instance is publicly exposed without an API key! It is highly recommended to restart with the \"--secure\" argument!")
+else:
+    print("No API key given because you are running locally.")
+
 
 def is_authorize_ignored(request):
     view_func = app.view_functions.get(request.endpoint)
@@ -418,6 +442,7 @@ def is_authorize_ignored(request):
             return True
     return False
 
+
 @app.before_request
 def before_request():
     # Request time measuring
@@ -426,14 +451,14 @@ def before_request():
     # Checks if an API key is present and valid, otherwise return unauthorized
     # The options check is required so CORS doesn't get angry
     try:
-        if request.method != 'OPTIONS' and is_authorize_ignored(request) == False and getattr(request.authorization, 'token', '') != api_key:
+        if request.method != 'OPTIONS' and args.secure and is_authorize_ignored(request) == False and getattr(request.authorization, 'token', '') != api_key:
             print(f"WARNING: Unauthorized API key access from {request.remote_addr}")
             response = jsonify({ 'error': '401: Invalid API key' })
             response.status_code = 401
-            return "this space is only for doctord98 but you can duplicate it and enjoy"
+            return response
     except Exception as e:
         print(f"API key check error: {e}")
-        return "this space is only for doctord98 but you can duplicate it and enjoy"
+        return "401 Unauthorized\n{}\n\n".format(e), 401
 
 
 @app.after_request
@@ -645,7 +670,7 @@ def tts_speakers():
     ]
     return jsonify(voices)
 
-
+# Added fix for Silero not working as new files were unable to be created if one already existed. - Rolyat 7/7/23
 @app.route("/api/tts/generate", methods=["POST"])
 @require_module("silero-tts")
 def tts_generate():
@@ -657,8 +682,15 @@ def tts_generate():
     # Remove asterisks
     voice["text"] = voice["text"].replace("*", "")
     try:
+        # Remove the destination file if it already exists
+        if os.path.exists('test.wav'):
+            os.remove('test.wav')
+
         audio = tts_service.generate(voice["speaker"], voice["text"])
-        return send_file(audio, mimetype="audio/x-wav")
+        audio_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.path.basename(audio))
+
+        os.rename(audio, audio_file_path)
+        return send_file(audio_file_path, mimetype="audio/x-wav")
     except Exception as e:
         print(e)
         abort(500, voice["speaker"])
@@ -743,8 +775,6 @@ def chromadb_purge():
 
     count = collection.count()
     collection.delete()
-    #Write deletion to persistent folder
-    chromadb_client.persist()
     print("ChromaDB embeddings deleted", count)
     return 'Ok', 200
 
@@ -768,6 +798,11 @@ def chromadb_query():
         name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
     )
 
+    if collection.count() == 0:
+        print(f"Queried empty/missing collection for {repr(data['chat_id'])}.")
+        return jsonify([])
+
+
     n_results = min(collection.count(), n_results)
     query_result = collection.query(
         query_texts=[data["query"]],
@@ -793,6 +828,69 @@ def chromadb_query():
 
     return jsonify(messages)
 
+@app.route("/api/chromadb/multiquery", methods=["POST"])
+@require_module("chromadb")
+def chromadb_multiquery():
+    data = request.get_json()
+    if "chat_list" not in data or not isinstance(data["chat_list"], list):
+        abort(400, '"chat_list" is required and should be a list')
+    if "query" not in data or not isinstance(data["query"], str):
+        abort(400, '"query" is required')
+
+    if "n_results" not in data or not isinstance(data["n_results"], int):
+        n_results = 1
+    else:
+        n_results = data["n_results"]
+
+    messages = []
+
+    for chat_id in data["chat_list"]:
+        if not isinstance(chat_id, str):
+            continue
+
+        try:
+            chat_id_md5 = hashlib.md5(chat_id.encode()).hexdigest()
+            collection = chromadb_client.get_collection(
+                name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+            )
+
+            # Skip this chat if the collection is empty
+            if collection.count() == 0:
+                continue
+
+            n_results_per_chat = min(collection.count(), n_results)
+            query_result = collection.query(
+                query_texts=[data["query"]],
+                n_results=n_results_per_chat,
+            )
+            documents = query_result["documents"][0]
+            ids = query_result["ids"][0]
+            metadatas = query_result["metadatas"][0]
+            distances = query_result["distances"][0]
+
+            chat_messages = [
+                {
+                    "id": ids[i],
+                    "date": metadatas[i]["date"],
+                    "role": metadatas[i]["role"],
+                    "meta": metadatas[i]["meta"],
+                    "content": documents[i],
+                    "distance": distances[i],
+                }
+                for i in range(len(ids))
+            ]
+
+            messages.extend(chat_messages)
+        except Exception as e:
+            print(e)
+
+    #remove duplicate msgs, filter down to the right number
+    seen = set()
+    messages = [d for d in messages if not (d['content'] in seen or seen.add(d['content']))]
+    messages = sorted(messages, key=lambda x: x['distance'])[0:n_results]
+
+    return jsonify(messages)
+
 
 @app.route("/api/chromadb/export", methods=["POST"])
 @require_module("chromadb")
@@ -802,9 +900,14 @@ def chromadb_export():
         abort(400, '"chat_id" is required')
 
     chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()
-    collection = chromadb_client.get_or_create_collection(
-        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
-    )
+    try:
+        collection = chromadb_client.get_collection(
+            name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+        )
+    except Exception as e:
+        print(e)
+        abort(400, "Chat collection not found in chromadb")
+
     collection_content = collection.get()
     documents = collection_content.get('documents', [])
     ids = collection_content.get('ids', [])
@@ -847,8 +950,27 @@ def chromadb_import():
 
 
     collection.upsert(documents=documents, metadatas=metadatas, ids=ids)
+    print(f"Imported {len(ids)} (total {collection.count()}) content entries into {repr(data['chat_id'])}")
 
     return jsonify({"count": len(ids)})
 
+
+if args.share:
+    from flask_cloudflared import _run_cloudflared
+    import inspect
+
+    sig = inspect.signature(_run_cloudflared)
+    sum = sum(
+        1
+        for param in sig.parameters.values()
+        if param.kind == param.POSITIONAL_OR_KEYWORD
+    )
+    if sum > 1:
+        metrics_port = randint(8100, 9000)
+        cloudflare = _run_cloudflared(port, metrics_port)
+    else:
+        cloudflare = _run_cloudflared(port)
+    print("Running on", cloudflare)
+
 ignore_auth.append(tts_play_sample)
 app.run(host=host, port=port)