%cd /kaggle/working/newer_project/
from modules.lightning import sdxl_lightning
from modules.audio_ldm import audio_generation
from modules.graph_plotter import plot_equation
from modules.search_web import search
from modules.vqa import moondream
from modules.grounding_sam import GroundingSAM
from modules.extra_functions import calculator, execute_code, none
from llm_modules.llm import llm_usage
from rw_modules.xtts_voice import xtts_model
from llm_modules.function_extract import FunctionExecutor
import json
import torch
from IPython.display import Image

# Loads up all models
grounding_sam = GroundingSAM("/kaggle/working/swint.py", "/kaggle/working/gdinot-1.8m-odvg.pth", "/kaggle/working/repvit_sam.pt")
#xtts = xtts_model()
llm_model = llm_usage("/kaggle/working/Meta-Llama-3-8B-Instruct-Q4_K_M.gguf")
sdxl_model = sdxl_lightning("/kaggle/working/sdxl_files")
#moondream_model = moondream()
search_engine = search()
#audio_gen = audio_generation()
function_executor = FunctionExecutor()


prompt = """Lets define 
x = y + 55
z = x + 5
y = 4
what can z be? Think step by step"""
personality = "a kind, helpful, entertaining, genious, smart, fun, assistant"
name = "Llama 3"
previous_prompts = []
previous_answers = []
prompt = f"""SYSTEM: {name} MUST follow the format. Start with action, always do observation, after observation, and then always start with Final_response
Adhere faithfully to the prescribed format, ensuring each new line commences with either an Action, an Observation, or a Final Response. Remain steadfast in character, infusing your contributions with depth and emotional resonance. Endeavor to craft eloquent, detailed responses imbued with passion and authenticity.
Utilize your intellect and imagination to enhance the quality of your contributions. You shall reform problems if they are not correctly formatted. Do not try to do stray from the format given.
User: {prompt}"""


torch.cuda.empty_cache()
def llm_speech(prompt, function_executer, personality, name, previous_prompts, previous_answers, functions='default'):
    speech_output = ""
    for out in llm_model.process_llm(prompt, function_executer=function_executer.get_function, functions=functions, personality=personality, name=name, previous_prompts=previous_prompts, previous_answers=previous_answers):
        token = out['token']
        if 'token' in out and token != None:
            yield token
        else:
            #yield token['output']
            yield out['output']
full_output = ""                       
for output in llm_speech(prompt, functions='default', function_executer=function_executor, personality=personality, name=name, previous_prompts=previous_prompts, previous_answers=previous_answers):
    if isinstance(output, dict) and output.get('display') != None:
        for url in output['display']['images']:
            display(Image(url))
        del output
    elif isinstance(output, str):
        print(output, end="")
        full_output += output

previous_prompts.append(prompt)
previous_answers.append(full_output)
torch.cuda.empty_cache()