Update generate_prompts.py

generate_prompts.py

@@ -3,60 +3,74 @@ import re
 
 def insert_description(sentence, character, description):
     """
-    Integrates the character and its description at the beginning of the sentence if the character is mentioned.
+    Integrates the character, its type, and its description within the sentence right after the character's name is mentioned.
+
     Parameters:
-        - sentence (str): The original sentence.
-        - character (str): The character to be described.
-        - description (str): The description of the character.
+        - sentence (str): The original sentence where the description is to be inserted.
+        - character (str): The character whose description is to be added.
+        - description (dict): The dictionary containing the character's type and descriptive words.
+
     Returns:
-        str: The sentence modified to include the character and description at the beginning.
+        str: The modified sentence with the character's description if the character is present; otherwise, returns the original sentence.
     """
-    # Inserts character and description at the beginning of the sentence if the character is found.
-    character = character.lower()
-    # Remove everything after the newline character
-    cleaned_description = re.sub(r'\n.*', '', description)
-    # Remove non-alphabetic characters and quotes from the description
-    cleaned_description = re.sub(r'[^a-zA-Z\s,]', '', cleaned_description).replace("'", '').replace('"', '')
-    # Check if the character appears in the sentence
-    if character in sentence.lower():
-        # Insert the character and its description at the beginning of the sentence
-        modified_sentence = f"{character}: {cleaned_description.strip()}. {sentence}"
-        return modified_sentence
-    else:
-        return sentence
+    character_lower = character.lower()
+
+    # Use regex to find and replace the character's name with the name plus the description
+    modified_sentence = re.sub(
+        fr"\b{character}\b",
+        fr"{character.capitalize()}{description}",
+        sentence,
+        flags=re.IGNORECASE
+    )
+    return modified_sentence
 
 def process_text(sentence, character_dict):
     """
     Enhances the given sentence by incorporating descriptions for each mentioned character.
+    Falls back to the original sentence if `character_dict` is empty.
+
     Parameters:
-        - sentence (str): The original sentence.
-        - character_dict (dict): Dictionary mapping characters to their descriptions.
+        - sentence (str): The original sentence to be processed.
+        - character_dict (dict): A dictionary mapping characters to their respective descriptions.
+
     Returns:
-        str: The sentence with integrated character descriptions.
+        str: The sentence modified to include character descriptions where applicable, or the original sentence.
     """
-    # Modifies sentences in the given text based on character descriptions.
-    modified_sentence = sentence  # Initialize with the original sentence
+    try:
+        # If character_dict is empty, return the original sentence
+        if not character_dict:
+            print("Character descriptions are empty, returning the original sentence.")
+            return sentence
+
+        # Start with the original sentence
+        modified_sentence = sentence
 
-    # Iterate through each character in the dictionary
-    for character, descriptions in character_dict.items():
-        for description in descriptions:
-            # Update the sentence with the character and its description
+        for character, description in character_dict.items():
+            # Insert description into the sentence where the character is mentioned
             modified_sentence = insert_description(modified_sentence, character, description)
-    return modified_sentence
+
+        print(f'modified_sentence: {modified_sentence}')
+        return modified_sentence
+
+    except Exception as e:
+        print(f"Error processing text: {e}. Returning original sentence.")
+        return sentence
 
 
-def generate_prompt(text, sentence_mapping, character_dict, selected_style):
+def generate_prompt(sentence, character_dict, selected_style):
     """
     Generates a prompt for image generation based on the selected style and input text.
+
     Parameters:
         - style (str): The chosen illustration style.
         - text (str): The input text for the illustration.
+
     Returns:
         tuple: A prompt string.
     """
     # Retrieve the enhanced sentence associated with the original text
-    enhanced_sentence = sentence_mapping.get(text, text)
-    image_descriptions = process_text(enhanced_sentence, character_dict)
+    image_descriptions = process_text(sentence, character_dict)
     # Define prompts and other parameters
-    prompt = f"Make an illustration in {selected_style} style from: {image_descriptions}"
+    prompt = f"Create an illustration in {selected_style} style from: {image_descriptions}"
     return prompt
+