121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
import sys
|
|
import resource
|
|
import gradio as gr
|
|
from ctypes_bindings import *
|
|
from model_class import *
|
|
from mesh_utils import *
|
|
|
|
# Set environment variables
|
|
#os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
|
|
#os.environ["GRADIO_SERVER_PORT"] = 8080
|
|
os.environ["RKLLM_LOG_LEVEL"] = "1"
|
|
# Set resource limit
|
|
resource.setrlimit(resource.RLIMIT_NOFILE, (102400, 102400))
|
|
|
|
history = []
|
|
|
|
if __name__ == "__main__":
|
|
# Helper function to define initializing model before class is declared
|
|
# Without this, you would need to initialize the class before you select the model
|
|
def initialize_model(model):
|
|
global rkllm_model
|
|
# Have to unload previous model in single-threaded mode
|
|
try:
|
|
rkllm_model.release()
|
|
except:
|
|
print("No model loaded! Continuing with initialization...")
|
|
# Initialize RKLLM model
|
|
init_msg = "=========INITIALIZING==========="
|
|
print(init_msg)
|
|
sys.stdout.flush()
|
|
rkllm_model = RKLLMLoaderClass(model=model)
|
|
model_init = f"RKLLM Model, {rkllm_model.model_name} has been initialized successfully!"
|
|
print(model_init)
|
|
complete_init = "=============================="
|
|
print(complete_init)
|
|
output = [[f"<h4 style=\"text-align:center;\">{model_init}\n</h4>", None]]
|
|
sys.stdout.flush()
|
|
return output
|
|
# Helper function to stream LLM output into the chat box
|
|
def get_RKLLM_output(message, history):
|
|
try:
|
|
yield from rkllm_model.get_RKLLM_output(message, history)
|
|
except RuntimeError as e:
|
|
print(f"ERROR: {e}")
|
|
return history
|
|
|
|
# Create a Gradio interface
|
|
with gr.Blocks(title="Chat with RKLLM") as chatRKLLM:
|
|
available_models = available_models()
|
|
gr.Markdown("<div align='center'><font size='10'> RKLLM Chat </font></div>")
|
|
with gr.Tabs():
|
|
with gr.TabItem("Select Model"):
|
|
model_dropdown = gr.Dropdown(choices=available_models, label="Select Model", value="None", allow_custom_value=True)
|
|
statusBox = gr.Chatbot(height=100)
|
|
model_dropdown.input(initialize_model, [model_dropdown], [statusBox])
|
|
with gr.TabItem("Txt2Txt"):
|
|
txt2txt = gr.ChatInterface(fn=get_RKLLM_output, type="messages")
|
|
txt2txt.chatbot.height = "70vh"
|
|
txt2txt.chatbot.resizeable = True
|
|
|
|
with gr.TabItem("Txt2HTML"):
|
|
# Поле для ввода текста
|
|
input_text = gr.Textbox(label="Введите ваш запрос", placeholder="Введите текст...", lines=3)
|
|
# Кнопка отправки запроса
|
|
submit_button = gr.Button("Отправить")
|
|
# Компонент для отображения HTML-ответа
|
|
output_html = gr.HTML(label="Ответ LLM")
|
|
def process_txt2txt(input_text):
|
|
html = input_text.replace("\t", "").replace("\n", "")
|
|
result = f"<iframe srcdoc='{html}' width='100%' height='1000'></iframe>"
|
|
return result
|
|
|
|
submit_button.click(
|
|
process_txt2txt,
|
|
inputs=input_text,
|
|
outputs=output_html
|
|
)
|
|
|
|
# with gr.TabItem("Txt2Mesh"):
|
|
# with gr.Row():
|
|
# with gr.Column(scale=2):
|
|
# txt2txt = gr.ChatInterface(fn=get_RKLLM_output, type="messages")
|
|
# txt2txt.chatbot.height = "70vh"
|
|
# txt2txt.chatbot.resizeable = True
|
|
# with gr.Column(scale=2):
|
|
# # Add the text box for 3D mesh input and button
|
|
# mesh_input = gr.Textbox(
|
|
# label="3D Mesh Input",
|
|
# placeholder="Paste your 3D mesh in OBJ format here...",
|
|
# lines=5,
|
|
# )
|
|
# visualize_button = gr.Button("Visualize 3D Mesh")
|
|
# output_model = gr.Model3D(
|
|
# label="3D Mesh Visualization",
|
|
# interactive=False,
|
|
# )
|
|
# # Link the button to the visualization function
|
|
# visualize_button.click(
|
|
# fn=apply_gradient_color,
|
|
# inputs=[mesh_input],
|
|
# outputs=[output_model]
|
|
# )
|
|
print("\nNo model loaded yet!\n")
|
|
|
|
|
|
# Enable the event queue system.
|
|
chatRKLLM.queue()
|
|
# Start the Gradio application.
|
|
chatRKLLM.launch(
|
|
server_name="0.0.0.0",
|
|
server_port=8080,
|
|
share=False,
|
|
show_error=True,
|
|
quiet=False
|
|
)
|
|
|
|
print("====================")
|
|
print("RKLLM model inference completed, releasing RKLLM model resources...")
|
|
rkllm_model.release()
|
|
print("====================")
|