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