import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

# Загружаем модель с параметрами для слабых систем
model_name = "AxisCommunity/OrionPaxAI_1.0V"

# Используем torch_dtype=torch.float16, чтобы уменьшить вес в памяти в 2 раза
# device_map="auto" поможет распределить нагрузку
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.float16, 
    device_map="auto",
    low_cpu_mem_usage=True
)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def predict(text):
    return pipe(text, max_new_tokens=50)[0]['generated_text']

demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch()