generate_propmts.py

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