Hugging Face's logo

Spaces:
Starchik1
/
ai2
Sleeping

App Files Community
Starchik1 commited on
Commit
728186c
·
verified ·
1 Parent(s): 25ac054

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +50 -6
main.py CHANGED
@@ -1,7 +1,8 @@
1
- from flask import Flask, render_template_string, request, redirect, url_for
2
  from huggingface_hub import snapshot_download
3
  from transformers import pipeline
4
  import os
 
5
 
6
  app = Flask(__name__)
7
 
@@ -13,14 +14,26 @@ AVAILABLE_MODELS = {
13
  "Mistral-7B": "mistralai/Mistral-7B-v0.1"
14
  }
15
 
16
- # Папка для моделей
17
  MODEL_DIR = "models"
18
  os.makedirs(MODEL_DIR, exist_ok=True)
19
 
20
- # Активная модель
21
  current_model = None
22
  pipe = None
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
  @app.route("/", methods=["GET", "POST"])
26
  def index():
@@ -31,15 +44,15 @@ def index():
31
  model_name = request.form["model"]
32
  repo_id = AVAILABLE_MODELS[model_name]
33
 
34
- # Скачиваем модель
35
- snapshot_download(repo_id, local_dir=os.path.join(MODEL_DIR, model_name))
 
36
  return redirect(url_for("index"))
37
 
38
  elif "use" in request.form:
39
  model_name = request.form["model"]
40
  repo_id = AVAILABLE_MODELS[model_name]
41
 
42
- # Загружаем модель
43
  current_model = model_name
44
  pipe = pipeline("text-generation", model=repo_id, device_map="auto")
45
  return redirect(url_for("index"))
@@ -56,6 +69,11 @@ def index():
56
  return render_template_string(TEMPLATE, models=AVAILABLE_MODELS, current=current_model, answer=None)
57
 
58
 
 
 
 
 
 
59
  TEMPLATE = """
60
  <!DOCTYPE html>
61
  <html>
@@ -67,6 +85,7 @@ TEMPLATE = """
67
  select, input[type=text], button { padding: 10px; margin-top: 10px; width: 100%; }
68
  button { background: #007BFF; color: white; border: none; cursor: pointer; }
69
  button:hover { background: #0056b3; }
 
70
  </style>
71
  </head>
72
  <body>
@@ -84,6 +103,7 @@ TEMPLATE = """
84
  <button type="submit" name="use">Use Model</button>
85
  </form>
86
  <p><b>Current model:</b> {{ current if current else "None" }}</p>
 
87
  </div>
88
 
89
  <div class="card">
@@ -98,6 +118,30 @@ TEMPLATE = """
98
  </div>
99
  {% endif %}
100
  </div>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
101
  </body>
102
  </html>
103
  """
 
1
+ from flask import Flask, render_template_string, request, redirect, url_for, jsonify
2
  from huggingface_hub import snapshot_download
3
  from transformers import pipeline
4
  import os
5
+ import threading
6
 
7
  app = Flask(__name__)
8
 
 
14
  "Mistral-7B": "mistralai/Mistral-7B-v0.1"
15
  }
16
 
 
17
  MODEL_DIR = "models"
18
  os.makedirs(MODEL_DIR, exist_ok=True)
19
 
 
20
  current_model = None
21
  pipe = None
22
 
23
+ # Прогресс загрузки
24
+ download_status = {"model": None, "status": "idle"} # idle | downloading | done
25
+
26
+
27
+ def download_model(model_name, repo_id):
28
+ global download_status
29
+ download_status = {"model": model_name, "status": "downloading"}
30
+
31
+ try:
32
+ snapshot_download(repo_id, local_dir=os.path.join(MODEL_DIR, model_name))
33
+ download_status = {"model": model_name, "status": "done"}
34
+ except Exception as e:
35
+ download_status = {"model": model_name, "status": f"error: {e}"}
36
+
37
 
38
  @app.route("/", methods=["GET", "POST"])
39
  def index():
 
44
  model_name = request.form["model"]
45
  repo_id = AVAILABLE_MODELS[model_name]
46
 
47
+ # Запускаем скачивание в отдельном потоке
48
+ thread = threading.Thread(target=download_model, args=(model_name, repo_id))
49
+ thread.start()
50
  return redirect(url_for("index"))
51
 
52
  elif "use" in request.form:
53
  model_name = request.form["model"]
54
  repo_id = AVAILABLE_MODELS[model_name]
55
 
 
56
  current_model = model_name
57
  pipe = pipeline("text-generation", model=repo_id, device_map="auto")
58
  return redirect(url_for("index"))
 
69
  return render_template_string(TEMPLATE, models=AVAILABLE_MODELS, current=current_model, answer=None)
70
 
71
 
72
+ @app.route("/progress")
73
+ def progress():
74
+ return jsonify(download_status)
75
+
76
+
77
  TEMPLATE = """
78
  <!DOCTYPE html>
79
  <html>
 
85
  select, input[type=text], button { padding: 10px; margin-top: 10px; width: 100%; }
86
  button { background: #007BFF; color: white; border: none; cursor: pointer; }
87
  button:hover { background: #0056b3; }
88
+ #progress-box { margin-top: 10px; padding: 10px; background: #eee; border-radius: 5px; display: none; }
89
  </style>
90
  </head>
91
  <body>
 
103
  <button type="submit" name="use">Use Model</button>
104
  </form>
105
  <p><b>Current model:</b> {{ current if current else "None" }}</p>
106
+ <div id="progress-box">Checking status...</div>
107
  </div>
108
 
109
  <div class="card">
 
118
  </div>
119
  {% endif %}
120
  </div>
121
+
122
+ <script>
123
+ function checkProgress() {
124
+ fetch("/progress")
125
+ .then(response => response.json())
126
+ .then(data => {
127
+ let box = document.getElementById("progress-box");
128
+ if (data.status === "idle") {
129
+ box.style.display = "none";
130
+ } else {
131
+ box.style.display = "block";
132
+ if (data.status === "downloading") {
133
+ box.innerHTML = "Downloading model: " + data.model + "...";
134
+ } else if (data.status === "done") {
135
+ box.innerHTML = "✅ Model " + data.model + " downloaded successfully!";
136
+ } else {
137
+ box.innerHTML = "⚠️ Error: " + data.status;
138
+ }
139
+ }
140
+ });
141
+ }
142
+
143
+ setInterval(checkProgress, 2000);
144
+ </script>
145
  </body>
146
  </html>
147
  """