fancyfeast commited on
Commit
8f7873a
1 Parent(s): 588df50

Add support for writing a custom prompt

Browse files
Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -158,7 +158,7 @@ image_adapter.to("cuda")
158
 
159
  @spaces.GPU()
160
  @torch.no_grad()
161
- def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str) -> tuple[str, str]:
162
  torch.cuda.empty_cache()
163
 
164
  # 'any' means no length specified
@@ -188,6 +188,9 @@ def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str
188
 
189
  # Add name, length, word_count
190
  prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)
 
 
 
191
 
192
  # For debugging
193
  print(f"Prompt: {prompt_str}")
@@ -317,13 +320,15 @@ with gr.Blocks() as demo:
317
  name_input = gr.Textbox(label="Person/Character Name (if applicable)")
318
  gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.")
319
 
 
 
320
  run_button = gr.Button("Caption")
321
 
322
  with gr.Column():
323
  output_prompt = gr.Textbox(label="Prompt that was used")
324
  output_caption = gr.Textbox(label="Caption")
325
 
326
- run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_length, extra_options, name_input], outputs=[output_prompt, output_caption])
327
 
328
 
329
  if __name__ == "__main__":
 
158
 
159
  @spaces.GPU()
160
  @torch.no_grad()
161
+ def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str) -> tuple[str, str]:
162
  torch.cuda.empty_cache()
163
 
164
  # 'any' means no length specified
 
188
 
189
  # Add name, length, word_count
190
  prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length)
191
+
192
+ if custom_prompt.strip() != "":
193
+ prompt_str = custom_prompt.strip()
194
 
195
  # For debugging
196
  print(f"Prompt: {prompt_str}")
 
320
  name_input = gr.Textbox(label="Person/Character Name (if applicable)")
321
  gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.")
322
 
323
+ custom_prompt = gr.Textbox(label="Custom Prompt (optional, will override all other settings)")
324
+
325
  run_button = gr.Button("Caption")
326
 
327
  with gr.Column():
328
  output_prompt = gr.Textbox(label="Prompt that was used")
329
  output_caption = gr.Textbox(label="Caption")
330
 
331
+ run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_length, extra_options, name_input, custom_prompt], outputs=[output_prompt, output_caption])
332
 
333
 
334
  if __name__ == "__main__":