Rooni commited on
Commit
118b6a4
·
verified ·
1 Parent(s): 4463bfc

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +161 -0
app.py ADDED
@@ -0,0 +1,161 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+ import requests
4
+ from huggingface_hub import InferenceClient
5
+ from flask import Flask, request, jsonify
6
+ from flask_cors import CORS
7
+ import json
8
+
9
+ app = Flask(__name__)
10
+ CORS(app)
11
+
12
+ def get_random_api_key():
13
+ keys = os.getenv("KEYS", "").split(",")
14
+ if keys and keys[0]:
15
+ return random.choice(keys).strip()
16
+ else:
17
+ raise ValueError("API keys not found. Please set the KEYS environment variable.")
18
+
19
+ def generate_story(prompt, style, stream=False):
20
+ try:
21
+ client = InferenceClient(api_key=get_random_api_key())
22
+
23
+ messages = [
24
+ {"role": "system", "content": f"Напиши хорошую историю в стиле '{style}'. Подробную, понятную, человечную (с душой), уникальную. Не обязательно делать концовку, можно только начало длинной истории. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
25
+ {"role": "user", "content": prompt}
26
+ ]
27
+
28
+ completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=stream)
29
+ if stream:
30
+ story = ""
31
+ for chunk in completion:
32
+ if chunk.choices:
33
+ story += chunk.choices[0].delta.content or ""
34
+ yield story
35
+ else:
36
+ story = completion.choices[0].message.content
37
+ yield story
38
+ except Exception as e:
39
+ yield f"Ошибка генерации: {e}"
40
+
41
+ def edit_story(original_story, edited_prompt, stream=False):
42
+ if not original_story:
43
+ yield f"Сначала сгенерируйте историю!"
44
+ return
45
+ if not edited_prompt:
46
+ yield f"Для начала введите что изменить."
47
+ return
48
+
49
+ try:
50
+ client = InferenceClient(api_key=get_random_api_key())
51
+ messages = [
52
+ {"role": "system", "content": "Отредактируй историю, учитывая предоставленные указания. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
53
+ {"role": "user", "content": edited_prompt},
54
+ {"role": "assistant", "content": original_story}
55
+ ]
56
+ completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=155000, stream=stream)
57
+ if stream:
58
+ edited_story = ""
59
+ for chunk in completion:
60
+ if chunk.choices:
61
+ edited_story += chunk.choices[0].delta.content or ""
62
+ yield edited_story
63
+ else:
64
+ edited_story = completion.choices[0].message.content
65
+ yield edited_story
66
+ except Exception as e:
67
+ yield f"Ошибка редактирования: {e}"
68
+
69
+ def next_story_func(original_story, next_prompt, continuation_type="Продолжение", stream=False):
70
+ if not original_story:
71
+ yield f"Сначала сгенерируйте историю!"
72
+ return
73
+
74
+ try:
75
+ client = InferenceClient(api_key=get_random_api_key())
76
+ if continuation_type == "Продолжение следует...":
77
+ continuation_prompt = "Продолжи, подведя к тому, что продолжение следует."
78
+ elif continuation_type == "Концовка":
79
+ continuation_prompt = "Подведи к концовке истории."
80
+ else:
81
+ continuation_prompt = next_prompt
82
+
83
+ messages = [
84
+ {"role": "system", "content": "Продли историю, учитывая предоставленные указания. Продливай В ТОЧНОСТИ С КОНЦА, прям с того же символа, слова, предложения. (В начале добавляй новые строки/пробел если надо для отступа). Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
85
+ {"role": "user", "content": continuation_prompt},
86
+ {"role": "assistant", "content": original_story}
87
+ ]
88
+ completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=stream)
89
+ if stream:
90
+ next_story = ""
91
+ for chunk in completion:
92
+ if chunk.choices:
93
+ next_story += chunk.choices[0].delta.content or ""
94
+ yield next_story
95
+ else:
96
+ next_story = completion.choices[0].message.content
97
+ yield next_story
98
+ except Exception as e:
99
+ yield f"Ошибка продления: {e}"
100
+
101
+ @app.route('/generate', methods=['POST'])
102
+ def api_generate_story():
103
+ data = request.get_json()
104
+ if not data or 'input' not in data:
105
+ return jsonify({"error": True, "message": "Missing 'input' in request body"}), 400
106
+
107
+ prompt = data['input']
108
+ style = data.get('style', 'Приключенческая') # Default style
109
+ stream = data.get('stream', False)
110
+
111
+ try:
112
+ story_generator = generate_story(prompt, style, stream)
113
+ if stream:
114
+ return jsonify({"error": False, "message": "".join(story_generator)})
115
+ else:
116
+ return jsonify({"error": False, "message": next(story_generator)})
117
+ except Exception as e:
118
+ return jsonify({"error": True, "message": str(e)}), 500
119
+
120
+ @app.route('/edit', methods=['POST'])
121
+ def api_edit_story():
122
+ data = request.get_json()
123
+ if not data or 'input' not in data or 'original' not in data:
124
+ return jsonify({"error": True, "message": "Missing 'input' or 'original' in request body"}), 400
125
+
126
+ original_story = data['original']
127
+ edited_prompt = data['input']
128
+ stream = data.get('stream', False)
129
+
130
+ try:
131
+ edited_story_generator = edit_story(original_story, edited_prompt, stream)
132
+ if stream:
133
+ return jsonify({"error": False, "message": "".join(edited_story_generator)})
134
+ else:
135
+ return jsonify({"error": False, "message": next(edited_story_generator)})
136
+ except Exception as e:
137
+ return jsonify({"error": True, "message": str(e)}), 500
138
+
139
+ @app.route('/continue', methods=['POST'])
140
+ def api_continue_story():
141
+ data = request.get_json()
142
+ if not data or 'input' not in data or 'original' not in data:
143
+ return jsonify({"error": True, "message": "Missing 'input' or 'original' in request body"}), 400
144
+
145
+ original_story = data['original']
146
+ next_prompt = data['input']
147
+ continuation_type = data.get('type', 'Продолжение') # Default continuation type
148
+ stream = data.get('stream', False)
149
+
150
+ try:
151
+ next_story_generator = next_story_func(original_story, next_prompt, continuation_type, stream)
152
+ if stream:
153
+ return jsonify({"error": False, "message": "".join(next_story_generator)})
154
+ else:
155
+ return jsonify({"error": False, "message": next(next_story_generator)})
156
+ except Exception as e:
157
+ return jsonify({"error": True, "message": str(e)}), 500
158
+
159
+ if __name__ == '__main__':
160
+ app.run(host='0.0.0.0', port=7860, debug=True)
161
+