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