utility/image_generator.py

import re
from PIL import Image
import io
from dotenv import load_dotenv
import os
import requests

load_dotenv()


def generate_image_prompts(script):
    # Split the script into sentences
    sentences = re.split(r'(?<=[.!?]) +', script)
    
    # Generate prompts for each sentence
    prompts = []
    for sentence in sentences:
        if sentence.strip():  # Ensure the sentence is not empty
            prompts.append(sentence.strip())
    
    return prompts

def hf_pipeline(prompt, max_retries=5, delay=30):
    retries = 0
    while retries < max_retries:
        response = requests.post(f"https://api-inference.huggingface.co/models/Shakker-Labs/AWPortrait-FL", 
                                 json={"inputs": prompt})
        if response.status_code == 503:
            print(f"Model is loading, retrying in {delay} seconds...")
            retries += 1
            time.sleep(delay)
        elif response.status_code == 200:
            return response.json()
        else:
            raise Exception(f"Failed to generate image. Status code: {response.status_code}, {response.text}")
    
    raise Exception(f"Failed to generate image after {max_retries} retries.")

def generate_images(prompts):
    try:
        image_files = []
        for idx, prompt in enumerate(prompts):
            print(f"Generating image for prompt: {prompt}")
            # Ensure the prompt is processed on the correct device
            image = hf_pipeline(prompt).images[0]
            filename = f"generated_image_{idx}.png"
            image.save(filename)
            image_files.append(filename)
    
        return image_files
    except Exception as e:
        print(e)