ai-meme-generator / ImageGenerator.py
islasher's picture
Update ImageGenerator.py
522336b verified
import os
# from PIL import Image
# from io import BytesIO
HF_API_KEY = os.getenv("HF_API_KEY")
# def generate_image(prompt):
# """Generate image using HF Inference API."""
# url = "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5" #"https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
# headers = {"Authorization": f"Bearer {HF_API_KEY}"}
# response = requests.post(url, headers=headers, json={"inputs": prompt})
# return Image.open(BytesIO(response.content))
import requests
from PIL import Image, UnidentifiedImageError, ImageDraw, ImageFont
from io import BytesIO
# def generate_image(prompt):
# """
# Generate an image from a prompt via an API.
# Returns a PIL.Image object. If the API fails, returns a placeholder image.
# """
# try:
# api_url="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"
# headers = {"Authorization": f"Bearer {HF_API_KEY}"}
# response = requests.post(api_url, headers=headers, json={"inputs": prompt.split('-')[-1]})
# # Check if the response is OK
# if response.status_code == 200:
# try:
# img = Image.open(BytesIO(response.content))
# return img
# except UnidentifiedImageError:
# print("API response is not a valid image.")
# else:
# print(f"API returned status {response.status_code}")
# print("Response content:", response.text[:200])
# except requests.RequestException as e:
# print(f"Request failed: {e}")
# # Fallback placeholder image
# print("Using placeholder image.")
# placeholder = Image.new("RGB", (512, 512), color=(200, 200, 200))
# draw = ImageDraw.Draw(placeholder)
# text = "Image\nNot\nAvailable"
# # Optional: Use a basic font
# try:
# font = ImageFont.load_default()
# w, h = draw.multiline_textsize(text, font=font)
# draw.multiline_text(
# ((512 - w) / 2, (512 - h) / 2), text, fill=(50, 50, 50), font=font, align="center"
# )
# except Exception:
# draw.text((50, 200), text, fill=(50, 50, 50))
# return placeholder
import os
import requests
from PIL import Image, UnidentifiedImageError, ImageDraw, ImageFont
from io import BytesIO
HF_API_KEY = os.getenv("HF_API_KEY")
def generate_imagee(prompt):
"""
Generate an image from the part of the prompt after the dash (-) via Hugging Face API.
Returns a PIL.Image object. If the API fails, returns a placeholder image.
"""
api_url = "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"
headers = {"Authorization": f"Bearer {HF_API_KEY}"}
# Get the part after the dash
if "-" in prompt:
prompt_to_send = prompt.split("-", 1)[-1].strip()
else:
prompt_to_send = prompt.strip()
try:
response = requests.post(api_url, headers=headers, json={"inputs": 'A meme of a person trying to rent a house, but the landlord is being super picky and only wants to rent to people who can afford to pay a lot of money for a small space.'})#prompt_to_send})
if response.status_code == 200:
try:
img = Image.open(BytesIO(response.content))
return img
except UnidentifiedImageError:
print("API response is not a valid image.")
else:
print(f"API returned status {response.status_code}")
print("Response content:", response.text[:200])
except requests.RequestException as e:
print(f"Request failed: {e}")
# Fallback placeholder image
print("Using placeholder image.")
placeholder = Image.new("RGB", (512, 512), color=(200, 200, 200))
draw = ImageDraw.Draw(placeholder)
text = "Image\nNot\nAvailable"
try:
font = ImageFont.load_default()
w, h = draw.multiline_textsize(text, font=font)
draw.multiline_text(
((512 - w) / 2, (512 - h) / 2), text, fill=(50, 50, 50), font=font, align="center"
)
except Exception:
draw.text((50, 200), text, fill=(50, 50, 50))
return placeholder
# -------------------
from diffusers import StableDiffusionPipeline
import torch
model_id = "sd-legacy/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
# pipe = pipe.to("cuda")
def generate_imagee2(prompt):
# Get the part after the dash
print('PROMPT TO INTRODUCE TO STABLDIFFUSION:', prompt)
prompt_to_send = extract_assistant_response(prompt)
print('INTRODUCED: ', prompt_to_send)
# if "-" in prompt:
# prompt_to_send = prompt.split("-", 1)[-1].strip()
# else:
# prompt_to_send = prompt.strip()
# prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt_to_send).images[0]
return image
def extract_assistant_response(decoded_text):
if "<|im_start|>assistant" in decoded_text:
part = decoded_text.split("<|im_start|>assistant", 1)[1]
part = part.split("<|im_end|>", 1)[0]
return part.strip()
return decoded_text.strip()
# image.save("astronaut_rides_horse.png")
# @InProceedings{Rombach_2022_CVPR,
# author = {Rombach, Robin and Blattmann, Andreas and Lorenz, Dominik and Esser, Patrick and Ommer, Bj\"orn},
# title = {High-Resolution Image Synthesis With Latent Diffusion Models},
# booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
# month = {June},
# year = {2022},
# pages = {10684-10695}
# }
import torch
from diffusers import FluxPipeline
# pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
# pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
# def generate_image(prompt):
# print('GENERATING IMAGE')
# prompt =extract_assistant_response(prompt)# "A cat holding a sign that says hello world"
# image = pipe(
# prompt,
# height=1024,
# width=1024,
# guidance_scale=3.5,
# num_inference_steps=50,
# max_sequence_length=512,
# generator=torch.Generator("cpu").manual_seed(0)
# ).images[0]
# #image.save("flux-dev.png")
# return image
from diffusers import AutoPipelineForText2Image
import torch
pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
# image = pipeline(prompt).images[0]
def generate_image(prompt):
print('GENERATING IMAGE')
prompt2 =extract_assistant_response(prompt)# "A cat holding a sign that says hello world"
print(prompt2)
image = pipe(prompt=prompt2, num_inference_steps=1, guidance_scale=0.0).images[0]
# image = pipe(
# prompt,
# height=1024,
# width=1024,
# guidance_scale=3.5,
# num_inference_steps=50,
# max_sequence_length=512,
# generator=torch.Generator("cpu").manual_seed(0)
# ).images[0]
#image.save("flux-dev.png")
return image