app.py

from fastapi import FastAPI, HTTPException, WebSocket, File, UploadFile, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import PlainTextResponse
from transformers import AutoTokenizer, AutoModelForCausalLM
from pydub import AudioSegment
import speech_recognition as sr
import io
import json
import torch
import datetime
from pydantic import BaseModel
import os
import sys
import pickle
import numpy
import librosa
import multiprocessing
import threading
import asyncio
# python3.10 -m uvicorn app:app --port 7860 --log-level warning
#  nvidia-smi
#  kill -9
path_to_add = os.path.join(os.path.dirname(__file__), "Wav2Lip")

if path_to_add not in sys.path:
    sys.path.insert(0, path_to_add)
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

active_websockets = []
from avatar import Avatar


class Item(BaseModel):
    text: str


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    active_websockets.append(websocket)
    global to_add_next_video
    global use_video_with_audio
    try:
        clientDisconnect = False
        while True:
            try:
                data = await websocket.receive_text()
                # print(f"Received event from client: {data}")
                event_data = json.loads(data)
                if event_data['event'] == "video-nearing-end":
                    to_add_next_video = True
                if event_data['event'] == "Button was clicked":
                    use_video_with_audio = True
            except WebSocketDisconnect:

                print("Client disconnected")
                clientDisconnect = True
                active_websockets.remove(websocket)
                break
    except Exception as e:
        print("Error:", e)
    finally:
        # print("in finally application_state ="+str(websocket.application_state))
        print("in finally clientDisconnect =" + str(clientDisconnect))

        # if not websocket.application_state == "disconnected":
        if not clientDisconnect:
            await websocket.close()
            print("WebSocket connection closed")


# can not have /submit-text/

@app.post("/submit-text")
async def submit_text(item: Item):
    global chat_history_ids
    global ts_file_names_and_duration
    input_text = item.text
    # new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
    #
    #
    #     # We assume that 'chat_history_ids' is already on the correct device and properly managed outside this snippet
    # # bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
    #
    #
    # # generated a response while limiting the total chat history to 1000 tokens,
    # # no_repeat_ngram_size=2,  early_stopping=True,  temperature=0.1,  top_p=0.9,
    # output_ids = lm_model.generate(new_user_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    # output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    # if output.startswith(input_text):
    #     output = output[len(input_text):].strip()
    # print("output=" + output)
    avatar.dir_clean_up()
    avatar.export_video = True
    # avatar.text_to_lip_video(user_input, inference_progress_bar)
    ts_file_names_and_duration = avatar.text_to_lip_video(input_text)
    return {"Response": input_text}


@app.get("/check-cuda")
async def check_cuda():
    if torch.cuda.is_available():
        print("cuda_ availabe")
        return {"CUDA available": True}
    else:
        print("cuda_ not availabe")
        return {"CUDA available": False}


@app.get("/hls/output.m3u8", response_class=PlainTextResponse)
async def generate_playlist():
    global to_add_next_video
    global use_video_with_audio
    global playlist_content
    global ts_file_names_and_duration
    global current_play_list_duration
    global time_to_speech
    current_time = datetime.datetime.now()
    formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
    # print("Formatted date and time:", formatted_time)
    # print(f"segment_counter={segment_counter} {formatted_time}")
    segment_duration = 6  # duration of the segment in seconds

    # playlist_content.append(f"#EXTINF:{segment_duration},")
    # playlist_content.append("output0.ts")

    # Add the segment multiple times based on the counter
    if to_add_next_video:

        # for _ in range(segment_counter + 1):
        # if  use_video_with_audio:
        #     playlist_content.append(f"#EXTINF:4.000000,")
        #     playlist_content.append("output_with_audio0.ts")
        #     playlist_content.append(f"#EXTINF:3.700000,")
        #     playlist_content.append("output_with_audio1.ts")
        #     playlist_content.append("#EXT-X-DISCONTINUITY")
        #     use_video_with_audio=False
        #     playlist_content.append(f"#EXTINF:6,")
        #     playlist_content.append("output_with_ambience0.ts")
        #     playlist_content.append("#EXT-X-DISCONTINUITY")
        if ts_file_names_and_duration:
            time_to_speech = current_play_list_duration
            for websocket in active_websockets:
                print(f"time to speech = {time_to_speech}")
                await websocket.send_text(f"time to speech = {time_to_speech}")
            print("ts_file_names_and_duration is not empty")
            # playlist_content.append(f"#EXTINF:4.000000,")
            # playlist_content.append("output_with_audio0.ts")
            # playlist_content.append(f"#EXTINF:3.700000,")
            # playlist_content.append("output_with_audio1.ts")
            playlist_content.append("#EXT-X-DISCONTINUITY")
            print(ts_file_names_and_duration)
            for file_name in ts_file_names_and_duration:
                print(file_name + " " + str(ts_file_names_and_duration[file_name]))
                playlist_content.append(f"#EXTINF:{ts_file_names_and_duration[file_name]},")
                current_play_list_duration = current_play_list_duration + float(ts_file_names_and_duration[file_name])
                playlist_content.append(f"{file_name}")
            playlist_content.append("#EXT-X-DISCONTINUITY")
            # playlist_content.append(f"#EXTINF:6,")
            # playlist_content.append("output_with_ambience0.ts")
            # playlist_content.append("#EXT-X-DISCONTINUITY")
            ts_file_names_and_duration = {}
        else:
            print("ts_file_names_and_duration is empty")
            playlist_content.append(f"#EXTINF:3,")
            playlist_content.append("output_3_seconds_with_ambience0.ts")
            playlist_content.append("#EXT-X-DISCONTINUITY")
            current_play_list_duration=current_play_list_duration+3
            # playlist_content.append(f"#EXTINF:6,")
            # playlist_content.append("output_with_ambience0.ts")
            # playlist_content.append("#EXT-X-DISCONTINUITY")
        to_add_next_video = False
    # print("\n".join(playlist_content))
    return "\n".join(playlist_content)


app.mount("/hls", StaticFiles(directory="hls"), name="hls")


def workhorse_function(conn, text, tokenizer, lm_model, avatar, device):
    print("workhorse_function, text=" + text)
    user_input_ids = tokenizer.encode(text + tokenizer.eos_token, return_tensors='pt').to(device)

    # We assume that 'chat_history_ids' is already on the correct device and properly managed outside this snippet
    # bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)

    # generated a response while limiting the total chat history to 1000 tokens,
    ai_result_ids = lm_model.generate(user_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    ai_output_text = tokenizer.decode(ai_result_ids[0], skip_special_tokens=True)
    if ai_output_text.startswith(text):
        ai_output_text = ai_output_text[len(text):].strip()
    print("output=" + ai_output_text)
    conn.send("step 4 done")
    avatar.dir_clean_up()
    avatar.export_video = True
    # avatar.text_to_lip_video(user_input, inference_progress_bar)
    ts_file_names_and_duration = avatar.text_to_lip_video(ai_output_text,conn)
    conn.send(ts_file_names_and_duration)

    # return ts_file_names_and_duration


def worker(conn):
    options = ['Aude', 'Kyla', 'Liv', 'Liv_3_seconds']
    images = ['ref_videos/Aude.png', 'ref_videos/Kyla.png', 'ref_videos/Liv.png']

    segment_counter = 0
    use_video_with_audio = False
    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
    lm_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    lm_model.to(device)
    init = False;
    input_text = "Hi, how are you?"
    new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
    if init:
        # We assume that 'chat_history_ids' is already on the correct device and properly managed outside this snippet
        bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
    else:
        bot_input_ids = new_user_input_ids
    init = True
    # generated a response while limiting the total chat history to 1000 tokens,
    chat_history_ids = lm_model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    print("output=" + output)

    avatar = Avatar()
    avatar.export_video = False
    print("load model")
    avatar.load_model("checkpoint/wav2lip_gan.pth")
    print("load model finished")
    avatar.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(avatar.device)
    avatar.output_audio_path = "audio/"
    avatar.output_audio_filename = "result.wav"
    avatar.temp_lip_video_no_voice_path = "temp/"
    avatar.temp_lip_video_no_voice_filename = "result.avi"
    avatar.output_video_path = "results/"
    avatar.output_video_name = "result_voice.mp4"
    selected_option = "Liv_3_seconds"
    avatar.ref_video_path_and_filename = f"ref_videos/{selected_option}.mp4"

    print("get video full frames")
    avatar.get_video_full_frames(avatar.ref_video_path_and_filename)
    print("get video full frames done")
    avatar.face_detect_batch_size = 16
    avatar.face_det_results_path_and_name = f'ref_videos/{selected_option}_face_det_result.pkl'
    # avatar.create_face_detection_results(avatar.video_full_frames,True)

    print("load face detection result")
    face_det_results_dict = {}
    for option in options:
        with open(f'ref_videos/{option}_face_det_result.pkl', 'rb') as file:
            face_det_results_dict[option] = pickle.load(file)
    print("load face detection result done")
    avatar.face_detect_img_results = face_det_results_dict[selected_option]

    avatar.export_video = False
    # avatar.text_to_lip_video(user_input, inference_progress_bar)
    # avatar.text_to_lip_video("Hi, how are you")
    while True:
        text = conn.recv()
        if text == "COMMAND_STOP":
            break
        # result = workhorse_function(conn,text,tokenizer, lm_model,avatar, device)
        result = workhorse_function(conn, text, tokenizer, lm_model, avatar, device)
        # conn.send(result)
    conn.close()


def result_handler(conn):
    global ts_file_names_and_duration
    while True:
        result = conn.recv()
        if result:
            print(f"Result received in main process not null")
            if isinstance(result, dict):
                print(f"Result received in main process: {result}")
                ts_file_names_and_duration = result
            elif isinstance(result, str):
                for websocket in active_websockets:
                    print(f"sending string "+ result)
                    asyncio.run(websocket.send_text(f"From Thread: {result}"))
@app.post("/upload")
async def upload_audio(audioFile: UploadFile = File(...)):
    global chat_history_ids
    global ts_file_names_and_duration
    print("in upload_audio")
    # Convert audio to the appropriate format
    content = await audioFile.read()
    audio = AudioSegment.from_file(io.BytesIO(content), format="webm")
    audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
    buffer = io.BytesIO()
    audio.export(buffer, format="wav")
    buffer.seek(0)
    for websocket in active_websockets:
        print(f"sending string " + "step 1 done")
        await websocket.send_text(f"From Thread: step 1 done")
    recognizer = sr.Recognizer()
    with sr.AudioFile(buffer) as source:
        # Read the entire audio file
        audio_data = recognizer.record(source)
    for websocket in active_websockets:
        print(f"sending string " + "step 2 done")
        await websocket.send_text(f"From Thread: step 2 done")
    try:
        # Recognize speech using Google Web Speech API
        text = recognizer.recognize_google(audio_data)
        parent_conn.send(text)
        for websocket in active_websockets:
            print(f"sending string " + "step 3 done")
            await websocket.send_text(f"From Thread: step 3 done")
        return {"transcript": text + "from localhost"}
    except sr.UnknownValueError:
        return {"error": "Google Speech Recognition could not understand audio"}
    except sr.RequestError as e:
        return {"error": f"Could not request results from Google Speech Recognition service; {e}"}
    return {"transcript": text}


print(__name__)
current_play_list_duration = float(0)
time_to_speech = float (0)
if __name__ == "app":
    print("Main process ID:", os.getpid())
    print("in __main__")
    ts_file_names_and_duration = {}

    # https://www.nvidia.com/Download/index.aspx download driver
    # device manager -> display adapter -> right click Nvdia properties, driver version, after install new driver, maybe just last digits change
    # Open PowerShell as Administrator:
    # dism.exe /online /enable-feature /featurename:Microsoft-Windows-Subsystem-Linux /all /norestart
    # dism.exe /online /enable-feature /featurename:VirtualMachinePlatform /all /norestart
    # reboot
    # microsoft store, search ubuntu, Ubuntu 22.04.3 LTS
    # start menu ubuntu: crate user pass
    # wsl -l -v
    # should see linux now
    # Configure Docker Desktop to Use WSL 2:
    #
    # After installation, right-click the Docker icon in the system tray and select 'Settings'.
    # Under the 'General' tab, ensure "Use the WSL 2 based engine" is checked.
    # Go to the 'Resources' -> 'WSL Integration' and enable integration for your Ubuntu distribution.
    # nvidia-smi
    # docker run --gpus all nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 nvidia-smi
    # --gpus all in run options may have to type one by one
    # /home/zmbfeng/huggingface_cache  /root/.cache/huggingface bind mounts
    # HF_HOME=/root/.cache/huggingface
    # add --log-level warning in unicorn run command to stop access logs

    # ffmpeg -i Liv.mp4 -profile:v baseline -level 3.0 -start_number 0 -hls_time 5 -hls_list_size 0 -f hls output.m3u8
    # HLS
    # file:///D:/PycharmProjects/ai_companion/ai_companion.html
    to_add_next_video = False
    playlist_content = [
        "#EXTM3U",
        "#EXT-X-VERSION:3",
        # "#EXT-X-TARGETDURATION:7",  # Adjust according to the segment length
        "#EXT-X-TARGETDURATION:3",  # Adjust according to the segment length
        "#EXT-X-MEDIA-SEQUENCE:0",
        "#EXTINF:3,",
        "output_3_seconds_with_ambience0.ts",
        "#EXT-X-DISCONTINUITY",
        # "#EXTINF:6,",
        # "output_with_ambience0.ts",
        # "#EXT-X-DISCONTINUITY",
        # "#EXTINF:2.666667,",
        # "ai_output0.ts",
        # "#EXT-X-DISCONTINUITY",
        # "#EXT-X-ENDLIST",
    ]
    current_play_list_duration=3
    parent_conn, child_conn = multiprocessing.Pipe()
    p = multiprocessing.Process(target=worker, args=(child_conn,))
    p.start()

    # Start a thread to handle results
    t = threading.Thread(target=result_handler, args=(parent_conn,))
    t.start()

#     command = [
#         "ffmpeg",
#         "-i", input_file,
#         "-profile:v", "baseline",
#         "-level", "3.0",
#         "-start_number", "0",
#         "-hls_time", "4",
#         "-hls_list_size", "0",
#         "-f", "hls",
#         output_playlist
#     ]
# ffmpeg -i Liv_with_audio.mp4 -codec:v libx264 -profile:v baseline -level 3.0 -g 120 -start_number 0 -hls_time 4 -hls_list_size 0 -f hls output_with_audio#


# ffmpeg -i Liv_with_audio.mp4 -profile:v baseline -level 3.0 -c:v libx264 -b:v 1500k -maxrate 1500k -bufsize 3000k -vf "scale=-2:720" -g 120 -hls_time 4 -hls_playlist_type vod -hls_segment_filename "video%03d.ts" video.m3u8
# #
# # # Encode audio
# ffmpeg -i Liv_with_audio.mp4 -g 120 -c:a aac -b:a 192k -hls_time 4 -hls_playlist_type vod -hls_segment_filename "audio%03d.ts" audio.m3u8