Update server.py

server.py

@@ -97,8 +97,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"
+port = 7860
+host = "0.0.0.0"
+
 summarization_model = (
     args.summarization_model
     if args.summarization_model
@@ -109,23 +110,6 @@ classification_model = (
     if args.classification_model
     else DEFAULT_CLASSIFICATION_MODEL
 )
-captioning_model = (
-    args.captioning_model if args.captioning_model else DEFAULT_CAPTIONING_MODEL
-)
-keyphrase_model = (
-    args.keyphrase_model if args.keyphrase_model else DEFAULT_KEYPHRASE_MODEL
-)
-prompt_model = args.prompt_model if args.prompt_model else DEFAULT_PROMPT_MODEL
-embedding_model = (
-    args.embedding_model if args.embedding_model else DEFAULT_EMBEDDING_MODEL
-)
-
-sd_use_remote = False if args.sd_model else True
-sd_model = args.sd_model if args.sd_model else DEFAULT_SD_MODEL
-sd_remote_host = args.sd_remote_host if args.sd_remote_host else DEFAULT_REMOTE_SD_HOST
-sd_remote_port = args.sd_remote_port if args.sd_remote_port else DEFAULT_REMOTE_SD_PORT
-sd_remote_ssl = args.sd_remote_ssl
-sd_remote_auth = args.sd_remote_auth
 
 modules = (
     args.enable_modules if args.enable_modules and len(args.enable_modules) > 0 else []
@@ -142,18 +126,6 @@ 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 == "cpu" else torch.float16
 
-if "caption" in modules:
-    print("Initializing an image captioning model...")
-    captioning_processor = AutoProcessor.from_pretrained(captioning_model)
-    if "blip" in captioning_model:
-        captioning_transformer = BlipForConditionalGeneration.from_pretrained(
-            captioning_model, torch_dtype=torch_dtype
-        ).to(device)
-    else:
-        captioning_transformer = AutoModelForCausalLM.from_pretrained(
-            captioning_model, torch_dtype=torch_dtype
-        ).to(device)
-
 if "summarize" in modules:
     print("Initializing a text summarization model...")
     summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model)
@@ -171,69 +143,6 @@ if "classify" in modules:
         torch_dtype=torch_dtype,
     )
 
-if "keywords" in modules:
-    print("Initializing a keyword extraction pipeline...")
-    import pipelines as pipelines
-
-    keyphrase_pipe = pipelines.KeyphraseExtractionPipeline(keyphrase_model)
-
-if "prompt" in modules:
-    print("Initializing a prompt generator")
-    gpt_tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
-    gpt_tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-    gpt_model = AutoModelForCausalLM.from_pretrained(prompt_model)
-    prompt_generator = pipeline(
-        "text-generation", model=gpt_model, tokenizer=gpt_tokenizer
-    )
-
-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"
-    )
-    sd_device = torch.device(sd_device_string)
-    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)
-    sd_pipe.safety_checker = lambda images, clip_input: (images, False)
-    sd_pipe.enable_attention_slicing()
-    # pipe.scheduler = KarrasVeScheduler.from_config(pipe.scheduler.config)
-    sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        sd_pipe.scheduler.config
-    )
-elif "sd" in modules and sd_use_remote:
-    print("Initializing Stable Diffusion connection")
-    try:
-        sd_remote = webuiapi.WebUIApi(
-            host=sd_remote_host, port=sd_remote_port, use_https=sd_remote_ssl
-        )
-        if sd_remote_auth:
-            username, password = sd_remote_auth.split(":")
-            sd_remote.set_auth(username, password)
-        sd_remote.util_wait_for_ready()
-    except Exception as e:
-        # remote sd from modules
-        print(
-            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}"
-        )
-        modules.remove("sd")
-
-if "tts" in modules:
-    if not os.path.exists(SILERO_SAMPLES_PATH):
-        os.makedirs(SILERO_SAMPLES_PATH)
-    print("Initializing Silero TTS server")
-    from silero_api_server import tts
-
-    tts_service = tts.SileroTtsService(SILERO_SAMPLES_PATH)
-    if len(os.listdir(SILERO_SAMPLES_PATH)) == 0:
-        print("Generating Silero TTS samples...")
-        tts_service.update_sample_text(SILERO_SAMPLE_TEXT)
-        tts_service.generate_samples()
-
 if "chromadb" in modules:
     print("Initializing ChromaDB")
     import chromadb
@@ -277,15 +186,6 @@ def classify_text(text: str) -> list:
     return sorted(output, key=lambda x: x["score"], reverse=True)
 
 
-def caption_image(raw_image: Image, max_new_tokens: int = 20) -> str:
-    inputs = captioning_processor(raw_image.convert("RGB"), return_tensors="pt").to(
-        device, torch_dtype
-    )
-    outputs = captioning_transformer.generate(**inputs, max_new_tokens=max_new_tokens)
-    caption = captioning_processor.decode(outputs[0], skip_special_tokens=True)
-    return caption
-
-
 def summarize_chunks(text: str, params: dict) -> str:
     try:
         return summarize(text, params)
@@ -331,70 +231,6 @@ def normalize_string(input: str) -> str:
     output = " ".join(unicodedata.normalize("NFKC", input).strip().split())
     return output
 
-
-def extract_keywords(text: str) -> list:
-    punctuation = "(){}[]\n\r<>"
-    trans = str.maketrans(punctuation, " " * len(punctuation))
-    text = text.translate(trans)
-    text = normalize_string(text)
-    return list(keyphrase_pipe(text))
-
-
-def generate_prompt(keywords: list, length: int = 100, num: int = 4) -> str:
-    prompt = ", ".join(keywords)
-    outs = prompt_generator(
-        prompt,
-        max_length=length,
-        num_return_sequences=num,
-        do_sample=True,
-        repetition_penalty=1.2,
-        temperature=0.7,
-        top_k=4,
-        early_stopping=True,
-    )
-    return [out["generated_text"] for out in outs]
-
-
-def generate_image(data: dict) -> Image:
-    prompt = normalize_string(f'{data["prompt_prefix"]} {data["prompt"]}')
-
-    if sd_use_remote:
-        image = sd_remote.txt2img(
-            prompt=prompt,
-            negative_prompt=data["negative_prompt"],
-            sampler_name=data["sampler"],
-            steps=data["steps"],
-            cfg_scale=data["scale"],
-            width=data["width"],
-            height=data["height"],
-            restore_faces=data["restore_faces"],
-            enable_hr=data["enable_hr"],
-            save_images=True,
-            send_images=True,
-            do_not_save_grid=False,
-            do_not_save_samples=False,
-        ).image
-    else:
-        image = sd_pipe(
-            prompt=prompt,
-            negative_prompt=data["negative_prompt"],
-            num_inference_steps=data["steps"],
-            guidance_scale=data["scale"],
-            width=data["width"],
-            height=data["height"],
-        ).images[0]
-
-    image.save("./debug.png")
-    return image
-
-
-def image_to_base64(image: Image, quality: int = 75) -> str:
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG", quality=quality)
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    return img_str
-
-
 @app.before_request
 # Request time measuring
 def before_request():
@@ -527,144 +363,10 @@ def api_prompt():
     print("Prompt output:", prompts, sep="\n")
     return jsonify({"prompts": prompts})
 
-
-@app.route("/api/image", methods=["POST"])
-@require_module("sd")
-def api_image():
-    required_fields = {
-        "prompt": str,
-    }
-
-    optional_fields = {
-        "steps": 30,
-        "scale": 6,
-        "sampler": "DDIM",
-        "width": 512,
-        "height": 512,
-        "restore_faces": False,
-        "enable_hr": False,
-        "prompt_prefix": PROMPT_PREFIX,
-        "negative_prompt": NEGATIVE_PROMPT,
-    }
-
-    data = request.get_json()
-
-    # Check required fields
-    for field, field_type in required_fields.items():
-        if field not in data or not isinstance(data[field], field_type):
-            abort(400, f'"{field}" is required')
-
-    # Set optional fields to default values if not provided
-    for field, default_value in optional_fields.items():
-        type_match = (
-            (int, float)
-            if isinstance(default_value, (int, float))
-            else type(default_value)
-        )
-        if field not in data or not isinstance(data[field], type_match):
-            data[field] = default_value
-
-    try:
-        print("SD inputs:", data, sep="\n")
-        image = generate_image(data)
-        base64image = image_to_base64(image, quality=90)
-        return jsonify({"image": base64image})
-    except RuntimeError as e:
-        abort(400, str(e))
-
-
-@app.route("/api/image/model", methods=["POST"])
-@require_module("sd")
-def api_image_model_set():
-    data = request.get_json()
-
-    if not sd_use_remote:
-        abort(400, "Changing model for local sd is not supported.")
-    if "model" not in data or not isinstance(data["model"], str):
-        abort(400, '"model" is required')
-
-    old_model = sd_remote.util_get_current_model()
-    sd_remote.util_set_model(data["model"], find_closest=False)
-    # sd_remote.util_set_model(data['model'])
-    sd_remote.util_wait_for_ready()
-    new_model = sd_remote.util_get_current_model()
-
-    return jsonify({"previous_model": old_model, "current_model": new_model})
-
-
-@app.route("/api/image/model", methods=["GET"])
-@require_module("sd")
-def api_image_model_get():
-    model = sd_model
-
-    if sd_use_remote:
-        model = sd_remote.util_get_current_model()
-
-    return jsonify({"model": model})
-
-
-@app.route("/api/image/models", methods=["GET"])
-@require_module("sd")
-def api_image_models():
-    models = [sd_model]
-
-    if sd_use_remote:
-        models = sd_remote.util_get_model_names()
-
-    return jsonify({"models": models})
-
-
-@app.route("/api/image/samplers", methods=["GET"])
-@require_module("sd")
-def api_image_samplers():
-    samplers = ["Euler a"]
-
-    if sd_use_remote:
-        samplers = [sampler["name"] for sampler in sd_remote.get_samplers()]
-
-    return jsonify({"samplers": samplers})
-
-
 @app.route("/api/modules", methods=["GET"])
 def get_modules():
     return jsonify({"modules": modules})
 
-
-@app.route("/api/tts/speakers", methods=["GET"])
-def tts_speakers():
-    voices = [
-        {
-            "name": speaker,
-            "voice_id": speaker,
-            "preview_url": f"{str(request.url_root)}api/tts/sample/{speaker}",
-        }
-        for speaker in tts_service.get_speakers()
-    ]
-    return jsonify(voices)
-
-
-@app.route("/api/tts/generate", methods=["POST"])
-def tts_generate():
-    voice = request.get_json()
-    if "text" not in voice or not isinstance(voice["text"], str):
-        abort(400, '"text" is required')
-    if "speaker" not in voice or not isinstance(voice["speaker"], str):
-        abort(400, '"speaker" is required')
-    # Remove asterisks
-    voice["text"] = voice["text"].replace("*", "")
-    try:
-        audio = tts_service.generate(voice["speaker"], voice["text"])
-        return send_file(audio, mimetype="audio/x-wav")
-    except Exception as e:
-        print(e)
-        abort(500, voice["speaker"])
-
-
-@app.route("/api/tts/sample/<speaker>", methods=["GET"])
-def tts_play_sample(speaker: str):
-    return send_from_directory(SILERO_SAMPLES_PATH, f"{speaker}.wav")
-
-
 @app.route("/api/chromadb", methods=["POST"])
 @require_module("chromadb")
 def chromadb_add_messages():
@@ -756,22 +458,4 @@ def chromadb_query():
 
     return jsonify(messages)
 
-
-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)
-
 app.run(host=host, port=port)