更新server.py

server.py

@@ -1,36 +1,31 @@
-from functools import wraps
-from flask import (
-    Flask,
-    jsonify,
-    request,
-    Response,
-    render_template_string,
-    abort,
-    send_from_directory,
-    send_file,
-)
-from flask_cors import CORS
-from flask_compress import Compress
-import markdown
 import argparse
-from transformers import AutoTokenizer, AutoProcessor, pipeline
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
-from transformers import BlipForConditionalGeneration
-import unicodedata
-import torch
-import time
-import os
+import base64
 import gc
-import sys
+import hashlib
+import os
 import secrets
-from PIL import Image
-import base64
+import sys
+import time
+import unicodedata
+from functools import wraps
 from io import BytesIO
 from random import randint
+
+import markdown
+import torch
 import webuiapi
-import hashlib
+from colorama import Fore, Style
+from colorama import init as colorama_init
+from flask import (Flask, Response, abort, jsonify, render_template_string,
+                   request, send_file, send_from_directory)
+from flask_compress import Compress
+from flask_cors import CORS
+from PIL import Image
+from transformers import (AutoModelForCausalLM, AutoModelForSeq2SeqLM,
+                          AutoProcessor, AutoTokenizer,
+                          BlipForConditionalGeneration, pipeline)
+
 from constants import *
-from colorama import Fore, Style, init as colorama_init
 
 colorama_init()
 
@@ -129,9 +124,9 @@ 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"
+# [HF, Huggingface] Set port to 7860, set host to remote.
+port = 7860
+host = "0.0.0.0"
 summarization_model = (
     args.summarization_model
     if args.summarization_model
@@ -222,8 +217,8 @@ if "summarize" in modules:
     ).to(device)
 
 if "sd" in modules and not sd_use_remote:
-    from diffusers import StableDiffusionPipeline
-    from diffusers import EulerAncestralDiscreteScheduler
+    from diffusers import (EulerAncestralDiscreteScheduler,
+                           StableDiffusionPipeline)
 
     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'
@@ -536,23 +531,10 @@ def image_to_base64(image: Image, quality: int = 75) -> str:
     img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
     return img_str
 
-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"{Fore.YELLOW}{Style.BRIGHT}Your API key is {api_key}{Style.RESET_ALL}")
-elif args.share and args.secure != True:
-    print(f"{Fore.RED}{Style.BRIGHT}WARNING: This instance is publicly exposed without an API key! It is highly recommended to restart with the \"--secure\" argument!{Style.RESET_ALL}")
-else:
-    print(f"{Fore.YELLOW}{Style.BRIGHT}No API key given because you are running locally.{Style.RESET_ALL}")
 
+ignore_auth = []
+# [HF, Huggingface] Get password instead of text file.
+api_key = os.environ.get("password")
 
 def is_authorize_ignored(request):
     view_func = app.view_functions.get(request.endpoint)
@@ -571,14 +553,16 @@ 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:
-            print(f"{Fore.RED}{Style.NORMAL}WARNING: Unauthorized API key access from {request.remote_addr}{Style.RESET_ALL}")
+        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}")
+            if request.method == 'POST':
+                print(f"Incoming POST request with {request.headers.get('Authorization')}")
             response = jsonify({ 'error': '401: Invalid API key' })
             response.status_code = 401
-            return response
+            return "https://(hf_name)-(space_name).hf.space/"
     except Exception as e:
         print(f"API key check error: {e}")
-        return "401 Unauthorized\n{}\n\n".format(e), 401
+        return "https://(hf_name)-(space_name).hf.space/"
 
 
 @app.after_request
@@ -1101,9 +1085,10 @@ def chromadb_import():
 
 
 if args.share:
-    from flask_cloudflared import _run_cloudflared
     import inspect
 
+    from flask_cloudflared import _run_cloudflared
+
     sig = inspect.signature(_run_cloudflared)
     sum = sum(
         1