ы

app.py

@@ -2,6 +2,7 @@ import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from accelerate import dispatch_model
 
 # Токен, если модель приватная
 hf_token = os.getenv("HF_TOKEN", None)
@@ -14,22 +15,23 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code=True,
     use_fast=True,
 )
-# Ensure offload directory exists before loading
-os.makedirs("offload", exist_ok=True)
 
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
     token=hf_token,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    trust_remote_code=True,
+)
+
+# Dispatch model to devices with offloading
+os.makedirs("offload", exist_ok=True)
+model = dispatch_model(
+    model,
     device_map="auto",
     offload_dir="offload",
     offload_state_dict=True,
-    low_cpu_mem_usage=True,
-    trust_remote_code=True,
 )
 
-
-
 # Создаём пайплайн
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)