Update server.py

server.py

@@ -21,7 +21,6 @@ import torch
 import time
 import os
 import gc
-import sys
 import secrets
 from PIL import Image
 import base64
@@ -34,9 +33,6 @@ 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):
@@ -44,16 +40,6 @@ 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(
@@ -70,8 +56,6 @@ 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(
@@ -82,10 +66,11 @@ 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")
@@ -120,8 +105,8 @@ parser.add_argument(
 
 args = parser.parse_args()
 
-port = args.port if args.port else 5100
-host = "0.0.0.0" if args.listen else "localhost"
+port = 7860
+host = "0.0.0.0"
 summarization_model = (
     args.summarization_model
     if args.summarization_model
@@ -157,16 +142,12 @@ if len(modules) == 0:
     print(f"Example: --enable-modules=caption,summarize{Style.RESET_ALL}")
 
 # Models init
-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_string = "cuda:0" if torch.cuda.is_available() and not args.cpu else "cpu"
 device = torch.device(device_string)
-torch_dtype = torch.float32 if device_string != cuda_device  else torch.float16
+torch_dtype = torch.float32 if device_string == "cpu" 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.{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.YELLOW}{Style.BRIGHT}torch-cuda is not supported on this device. Defaulting to CPU mode.{Style.RESET_ALL}")
 
 print(f"{Fore.GREEN}{Style.BRIGHT}Using torch device: {device_string}{Style.RESET_ALL}")
 
@@ -203,10 +184,12 @@ 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_device if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+    print("Initializing Stable Diffusion pipeline")
+    sd_device_string = (
+        "cuda" if torch.cuda.is_available() and not args.sd_cpu else "cpu"
+    )
     sd_device = torch.device(sd_device_string)
-    sd_torch_dtype = torch.float32 if sd_device_string != cuda_device else torch.float16
+    sd_torch_dtype = torch.float32 if sd_device_string == "cpu" else torch.float16
     sd_pipe = StableDiffusionPipeline.from_pretrained(
         sd_model, custom_pipeline="lpw_stable_diffusion", torch_dtype=sd_torch_dtype
     ).to(sd_device)
@@ -269,19 +252,26 @@ if "chromadb" in modules:
     posthog.capture = lambda *args, **kwargs: None
     if args.chroma_host is None:
         if args.chroma_persist:
-            chromadb_client = chromadb.PersistentClient(path=args.chroma_folder, settings=Settings(anonymized_telemetry=False))
+            chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False, persist_directory=args.chroma_folder, chroma_db_impl='duckdb+parquet'))
             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.EphemeralClient(Settings(anonymized_telemetry=False))
+            chromadb_client = chromadb.Client(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.HttpClient(host=args.chroma_host, port=chroma_port, settings=Settings(anonymized_telemetry=False))
+        chromadb_client = chromadb.Client(
+            Settings(
+                anonymized_telemetry=False,
+                chroma_api_impl="rest",
+                chroma_server_host=args.chroma_host,
+                chroma_server_http_port=chroma_port
+            )
+        )
         print(f"ChromaDB is remotely configured at {args.chroma_host}:{chroma_port}")
 
-    chromadb_embedder = SentenceTransformer(embedding_model, device=device_string)
+    chromadb_embedder = SentenceTransformer(embedding_model)
     chromadb_embed_fn = lambda *args, **kwargs: chromadb_embedder.encode(*args, **kwargs).tolist()
 
     # Check if the db is connected and running, otherwise tell the user
@@ -415,24 +405,10 @@ 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 = []    
 
-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.")
-
+api_key = os.environ.get("password")
 
 def is_authorize_ignored(request):
     view_func = app.view_functions.get(request.endpoint)
@@ -442,7 +418,6 @@ def is_authorize_ignored(request):
             return True
     return False
 
-
 @app.before_request
 def before_request():
     # Request time measuring
@@ -451,14 +426,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 args.secure and is_authorize_ignored(request) == False and getattr(request.authorization, 'token', '') != api_key:
+        if request.method != 'OPTIONS' 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 response
+            return "this space is only for doctord98 but you can duplicate it and enjoy"
     except Exception as e:
         print(f"API key check error: {e}")
-        return "401 Unauthorized\n{}\n\n".format(e), 401
+        return "this space is only for doctord98 but you can duplicate it and enjoy"
 
 
 @app.after_request
@@ -670,7 +645,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():
@@ -682,15 +657,8 @@ 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"])
-        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")
+        return send_file(audio, mimetype="audio/x-wav")
     except Exception as e:
         print(e)
         abort(500, voice["speaker"])
@@ -775,6 +743,8 @@ def chromadb_purge():
 
     count = collection.count()
     collection.delete()
+    #Write deletion to persistent folder
+    chromadb_client.persist()
     print("ChromaDB embeddings deleted", count)
     return 'Ok', 200
 
@@ -798,11 +768,6 @@ 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"]],
@@ -828,69 +793,6 @@ 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")
@@ -900,14 +802,9 @@ def chromadb_export():
         abort(400, '"chat_id" is required')
 
     chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()
-    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 = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
     collection_content = collection.get()
     documents = collection_content.get('documents', [])
     ids = collection_content.get('ids', [])
@@ -950,27 +847,8 @@ 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)
+app.run(host=host, port=port)