Update app.py

app.py

@@ -53,15 +53,22 @@ def edit_story(original_story="", edited_prompt=""):
     except Exception as e:
         return f"Ошибка редактирования: {e}"
 
-def next_story_func(original_story="", next_prompt=""):
+def next_story_func(original_story="", next_prompt="", continuation_type="Продолжение"):
     if original_story == "":
         return f"Сначала сгенерируйте историю!"
     
     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": "Продли историю, учитывая предоставленные указания. Продливай В ТОЧНОСТИ С КОНЦА, прям с того же символа, слова, предложения. (В начале добавляй новые строки/пробел если надо для отступа)"},
-            {"role": "user", "content": next_prompt},
+            {"role": "user", "content": continuation_prompt},
             {"role": "assistant", "content": original_story}
         ]
         completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1000)
@@ -102,6 +109,7 @@ with gr.Blocks(css=css) as demo:
                     edone_button = gr.Button("Принять")
 
                 with gr.Tab("Продление"):
+                    continuation_type = gr.Radio(choices=["Продолжение", "Продолжение следует...", "Концовка"], label="Выберите тип продолжения", value="Продолжение")
                     next_prompt = gr.Textbox(label="Введите изменения для продления истории (Необязательно)", placeholder="Продолжи, но что бы было...", lines=5)
                     next_button = gr.Button("Продлить", variant='primary')
                     next_story_output = gr.Textbox(label="Продолжение истории", lines=10, placeholder="Здесь будет продолжение вашей истории!")
@@ -109,10 +117,10 @@ with gr.Blocks(css=css) as demo:
 
     generate_button.click(generate_story, inputs=[prompt, style], outputs=[output_story], concurrency_limit=250)
     edit_button.click(edit_story, inputs=[output_story, edited_prompt], outputs=[edited_story], concurrency_limit=250)
-    next_button.click(next_story_func, inputs=[output_story, next_prompt], outputs=[next_story_output], concurrency_limit=250)
+    next_button.click(next_story_func, inputs=[output_story, next_prompt, continuation_type], outputs=[next_story_output], concurrency_limit=250)
     
     edone_button.click(edone_story, inputs=[output_story, edited_story], outputs=[output_story], concurrency_limit=550)
     
     ndone_button.click(ndone_story, inputs=[output_story, next_story_output], outputs=[output_story], concurrency_limit=550)
 
-demo.launch(show_api=False, share=False)
+demo.launch(show_api=False, share=False)