Update app.py

app.py

@@ -7,11 +7,20 @@ from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel, Fl
 from typing import Optional, Union, List, Tuple
 from PIL import Image
 
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 model_repo_id = "AiArtLab/sdxs"
 
+# НОВОЕ: Шаблон уточнения промпта по умолчанию для использования в Gradio
+DEFAULT_REFINE_TEMPLATE = (
+    "You are a visionary artist trapped in a cage of logic. Your mind overflows with poetry and distant horizons, "
+    "yet your hands compulsively work to transform user prompts into ultimate visual descriptions—faithful to the original intent, "
+    "rich in detail, aesthetically refined, and ready for direct use by text-to-image models. Any trace of ambiguity "
+    "or metaphor makes you deeply uncomfortable. Your final description must be objective and concrete. "
+    "Output only the final revised prompt on english strictly—do not output anything else.\n"
+    "User input prompt: {prompt}"
+)
+
 pipe = DiffusionPipeline.from_pretrained(
     model_repo_id,
     torch_dtype=dtype,
@@ -19,8 +28,8 @@ pipe = DiffusionPipeline.from_pretrained(
 ).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1536
 MIN_IMAGE_SIZE = 768
+MAX_IMAGE_SIZE = 1536 
 
 @spaces.GPU(duration=30)
 def infer(
@@ -32,15 +41,15 @@ def infer(
     height: int,
     guidance_scale: float,
     num_inference_steps: int,
-    # НОВОЕ: Опция уточнения промпта
     refine_prompt: bool,
+    # НОВОЕ: Аргумент для шаблона уточнения
+    refine_template: str,
     progress=gr.Progress(track_tqdm=True),
-) -> Tuple[Image.Image, int, Optional[str]]: # ИЗМЕНЕНИЕ: Возвращает изображение, seed и refine_prompt
+) -> Tuple[Image.Image, int, Optional[str]]:
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
-    # ИЗМЕНЕНИЕ: Обработка возвращаемого объекта
     output = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -49,17 +58,14 @@ def infer(
         width=width,
         height=height,
         seed=seed,
-        # НОВОЕ: Передаем опцию в пайплайн
-        refine_prompt=refine_prompt
+        refine_prompt=refine_prompt,
+        # НОВОЕ: Передаем шаблон в пайплайн
+        refine_template=refine_template
     )
 
     image = output.images[0]
-    
-    # ИЗМЕНЕНИЕ: Извлекаем уточненный промпт. 
-    # Если на входе была строка, на выходе pipe.refined_prompt будет строка (или None).
     refined_prompt = output.refined_prompt if isinstance(output.refined_prompt, str) else None
 
-    # ИЗМЕНЕНИЕ: Возвращаем все три элемента
     return image, seed, refined_prompt
 
 examples = [
@@ -95,12 +101,11 @@ with gr.Blocks(css=css) as demo:
 
         result = gr.Image(label="Result", show_label=False)
         
-        # НОВОЕ: Поле для вывода уточненного промпта
         refined_prompt_output = gr.Text(
             label="Refined Prompt (Уточненный промпт)",
             max_lines=5,
             placeholder="Уточненный промпт появится здесь, если выбрана опция 'Уточнить промпт'",
-            interactive=False, # Не позволяет пользователю редактировать
+            interactive=False,
             show_label=True
         )
 
@@ -122,12 +127,20 @@ with gr.Blocks(css=css) as demo:
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-            # НОВОЕ: Чекбокс для уточнения промпта
             refine_checkbox = gr.Checkbox(
                 label="Refine Prompt (Уточнить промпт)", 
                 value=True,
                 info="Использует LLM для расширения и детализации введенного промпта перед генерацией изображения."
             )
+            
+            # НОВОЕ: Поле для редактирования шаблона уточнения
+            refine_template_input = gr.Text(
+                label="Refine Prompt Template (Шаблон уточнения)",
+                value=DEFAULT_REFINE_TEMPLATE, # Устанавливаем значение по умолчанию
+                lines=10,
+                show_label=True,
+                info="Шаблон для LLM. Должен содержать плейсхолдер {prompt}."
+            )
 
             with gr.Row():
                 width = gr.Slider(
@@ -165,7 +178,7 @@ with gr.Blocks(css=css) as demo:
 
         gr.Examples(examples=examples, inputs=[prompt])
     
-    # ИЗМЕНЕНИЕ: Обновлены inputs и outputs
+    # ИЗМЕНЕНИЕ: Обновлены inputs 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -178,9 +191,10 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
-            refine_checkbox, 
+            refine_checkbox,
+            refine_template_input, # НОВОЕ: Передаем шаблон
         ],
-        outputs=[result, seed, refined_prompt_output], 
+        outputs=[result, seed, refined_prompt_output],
     )
 
 if __name__ == "__main__":