Upload fl2basepromptgen.py

fl2basepromptgen.py

@@ -58,7 +58,8 @@ def fl_run_example(image):
     )
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = fl_processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
-    return fl_modify_caption(parsed_answer["<DESCRIPTION>"])
+    print(parsed_answer)
+    return fl_modify_caption(parsed_answer["<PROMPT>"])
 
 
 def predict_tags_fl2_base_prompt_gen(image: Image.Image, input_tags: str, algo: list[str]):