Create generate_propmts.py

generate_propmts.py

@@ -0,0 +1,105 @@
+import base64
+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.
+
+    Parameters:
+        - sentence (str): The original sentence.
+        - character (str): The character to be described.
+        - description (str): The description of the character.
+
+    Returns:
+        str: The sentence modified to include the character and description at the beginning.
+    """
+    # 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
+
+def process_text(sentence, character_dict):
+    """
+    Enhances the given sentence by incorporating descriptions for each mentioned character.
+
+    Parameters:
+        - sentence (str): The original sentence.
+        - character_dict (dict): Dictionary mapping characters to their descriptions.
+
+    Returns:
+        str: The sentence with integrated character descriptions.
+    """
+    # Modifies sentences in the given text based on character descriptions.
+    modified_sentence = sentence  # Initialize with the original 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
+            modified_sentence = insert_description(modified_sentence, character, description)
+    return modified_sentence
+
+
+def generate_prompt(text, sentence_mapping, character_dict, selected_style):
+    """
+    Generates a prompt and negative 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 tuple containing the prompt and negative prompt strings.
+    """
+    # Retrieve the enhanced sentence associated with the original text
+    enhanced_sentence = sentence_mapping.get(text, text)
+    image_descriptions = process_text(enhanced_sentence, character_dict)
+    # Define prompts and other parameters
+    prompt = f"Make an illustration in {selected_style} style from: {image_descriptions}"
+    negative_prompt = "lowres, bad anatomy, bad hands, text, chat box, words, error, missing fingers, extra digit, " \
+                      "fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, " \
+                      "watermark, username, blurry "
+    return prompt, negative_prompt
+
+def get_image_from_space(text, sentence_mapping, character_dict, selected_style, client):
+    """
+    Requests an image from a Hugging Face space based on the provided prompt.
+
+    Parameters:
+        - prompt (str): The text prompt for image generation.
+        - negative_prompt (str): Text specifying what to avoid in the image.
+
+    Returns:
+        bytes: The generated image data in bytes format, or None if the request fails.
+    """
+    image_bytes = None  # Initialize image bytes to None for error handling
+    try:
+        with st.spinner("עוד כמה רגעים והאיור יופיע"):
+            # Define the payload with the text prompt
+            prompt,_ = generate_prompt(text, sentence_mapping, character_dict, selected_style)
+            payload = {
+                "inputs": prompt
+            }
+            result = client.predict(
+                prompt=payload,
+                api_name="/predict"
+            )
+
+            # Check if the result is a base64 encoded string
+            if isinstance(result, str):
+                image_bytes = base64.b64decode(result)
+                return image_bytes
+            else:
+                st.error("אוי לא ניתן לייצר תמונה, יש לנסות שוב בהמשך")
+    except Exception as e:
+        print(f"אוי לא ניתן לייצר תמונה, יש לנסות שוב בהמשך: {e}")
+        return image_bytes