Spaces:
Running
Running
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) | |