Spaces:
Running
Running

Refactor image handling by removing debug print statements, adding image cleanup, and improving post creation sleep duration
9194514
import base64 | |
import os | |
import pollinations as ai | |
from PIL import Image | |
import io | |
import re | |
import requests | |
def extract_summary(text): | |
text = text.replace("#", "").strip().lower() | |
match = re.search(r"summary(.*?)highlights", text.replace("#", ""), re.DOTALL) | |
if match: | |
return match.group(1).replace("#", "").replace("\n", "").strip() | |
return None | |
def generate_image(title, summary): | |
try: | |
extracted_summary = extract_summary(summary) | |
if not extracted_summary: | |
extracted_summary = summary | |
data = r"{}: {}".format(title.replace("{", "").replace("}","").replace("\n", ""), extracted_summary).strip().replace("&", "&").replace("{", "").replace("}","").replace("\n", "") | |
model_obj = ai.ImageModel( | |
model=ai.turbo, | |
width=1280, | |
height=720, | |
seed=2342342340, | |
enhance=True, | |
nologo=True | |
) | |
model_obj.generate( | |
prompt=r"{}".format(data), | |
negative=("a bad image, low res, low quality, pixelated, blurry, bad lighting, bad angle, bad composition, bad focus, bad exposure, bad white balance, bad color, bad contrast, bad saturation, bad sharpness, bad noise, bad artifacts, bad framing, bad cropping, bad editing, bad filter, bad lens, bad camera, bad sensor, bad settings, bad post-processing, bad retouching, bad enhancement, bad manipulation, bad distortion, bad watermark, bad logo, bad text, bad overlay, bad background"), | |
save=True, | |
file="image.png" | |
) | |
with open("image.png", "rb") as image_file: | |
image_data = image_file.read() | |
base64_encoded = base64.b64encode(image_data).decode('utf-8') | |
data_uri = f"data:image/png;base64,{base64_encoded}" | |
return r"{}".format(data_uri).strip() | |
except Exception as e: | |
return f"An error occurred: {e}" | |
def verify_image(image_data): | |
try: | |
image_stream = io.BytesIO(image_data) | |
image = Image.open(image_stream) | |
image.verify() | |
image_stream.seek(0) | |
image = Image.open(image_stream) | |
return True | |
except Exception as e: | |
print(f"Error verifying image: {e}") | |
return False | |
def upload_image(data_uri, api_key): | |
try: | |
base64_image = data_uri.split(",")[1] | |
url = f'https://api.imgbb.com/1/upload?key={api_key}' | |
response = requests.post(url, data={'image': base64_image}).json() | |
if response.get('status') == 200: | |
return response['data']['url'] | |
else: | |
return "https://i.ibb.co/qdC1nSx/landscape-placeholder-e1608289113759.png" | |
except Exception as e: | |
print(f"Error uploading image: {e}") | |
return "https://i.ibb.co/qdC1nSx/landscape-placeholder-e1608289113759.png" | |
def fetch_image(title, summary, api_key): | |
image_url = "https://i.ibb.co/qdC1nSx/landscape-placeholder-e1608289113759.png" | |
data_uri = generate_image(title, summary) | |
image_data = base64.b64decode(data_uri.split(",")[1]) | |
if verify_image(image_data): | |
image_url = upload_image(data_uri, api_key) | |
os.remove("image.png") | |
return image_url | |