Initial mirror of VectorSpaceLab/EditScore@4609c5d

da50cbf verified about 2 months ago

7.07 kB

import base64
import requests
from io import BytesIO, StringIO
from typing import Union, Optional, Tuple, List
from PIL import Image, ImageOps
import os

def get_api_key(file_path):
    # Read the API key from the first line of the file
    with open(file_path, 'r') as file:
        return file.readline().strip()

# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def pick_next_item(current_item, item_list):
    if current_item not in item_list:
        raise ValueError("Current item is not in the list")
    current_index = item_list.index(current_item)
    next_index = (current_index + 1) % len(item_list)

    return item_list[next_index]

# Function to encode a PIL image
def encode_pil_image(pil_image):
    # Create an in-memory binary stream
    image_stream = BytesIO()
    
    # Save the PIL image to the binary stream in JPEG format (you can change the format if needed)
    pil_image.save(image_stream, format='JPEG')
    
    # Get the binary data from the stream and encode it as base64
    image_data = image_stream.getvalue()
    base64_image = base64.b64encode(image_data).decode('utf-8')
    
    return base64_image


def load_image(image: Union[str, Image.Image], format: str = "RGB", size: Optional[Tuple] = None) -> Image.Image:
    """
    Load an image from a given path or URL and convert it to a PIL Image.
    Args:
        image (Union[str, Image.Image]): The image path, URL, or a PIL Image object to be loaded.
        format (str, optional): Desired color format of the resulting image. Defaults to "RGB".
        size (Optional[Tuple], optional): Desired size for resizing the image. Defaults to None.
    Returns:
        Image.Image: A PIL Image in the specified format and size.
    Raises:
        ValueError: If the provided image format is not recognized.
    """
    if isinstance(image, str):
        if image.startswith("http://") or image.startswith("https://"):
            image = Image.open(requests.get(image, stream=True).raw)
        elif os.path.isfile(image):
            image = Image.open(image)
        else:
            raise ValueError(
                f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path"
            )
    elif isinstance(image, Image.Image):
        image = image
    else:
        raise ValueError(
            "Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image."
        )
    image = ImageOps.exif_transpose(image)
    image = image.convert(format)
    if (size != None):
        image = image.resize(size, Image.LANCZOS)
    return image

class GPT4v():
    def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4-vision-preview"):
        """OpenAI GPT-4-vision model wrapper
        Args:
            api_key_path (str): Path to the API key file. Defaults to 'keys/secret.env'.
            are_images_encoded (bool): Whether the images are encoded in base64. Defaults to False.
        """
        self.multiple_api_keys = False
        self.current_key_file = None
        self.api_key = key

        self.url = url
        self.model_name = model_name
        self.use_encode = are_images_encoded

    def prepare_input(self, image_links: List = [], text_prompt: str = ""):
        prompt_content = []
        text_dict = {
                    "type": "text",
                    "text": text_prompt
                }
        prompt_content.append(text_dict)

        if not isinstance(image_links, list):
            image_links = [image_links]
            
        for image_link in image_links:
            image = load_image(image_link)
            if self.use_encode:
                visual_dict = {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/jpeg;base64,{encode_pil_image(image)}"}
                    }
            else:
                visual_dict = {
                        "type": "image_url",
                        "image_url": {"url": image_link}
                    }
            prompt_content.append(visual_dict)
        return prompt_content

    def inference(self, prompt, seed: Optional[int] = None):
        payload = {
            "model": self.model_name,
            "messages": [
            {
                "role": "user",
                "content": prompt
            }
            ],
            # "max_tokens": 1400
        }
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        # try:
        response = requests.post(self.url, json=payload, headers=headers, timeout=180)  # Set timeout to 5 minutes (300 seconds)
        # except Exception as e:
        #     print(f"Error: {e}")
        #     return ""
        #return response.text
        return self.extract_response(response)
    
    def extract_response(self, response):
        try:
            response = response.json()
            out = response['choices'][0]['message']['content']
            return out
        except Exception as e:
            print(f"Error: {e}")
            if response['error']['code'] == 'content_policy_violation':
                print("Code is content_policy_violation")
            elif response['error']['code'] in ['rate_limit_exceeded', 'insufficient_quota', 'insufficient_user_quota']:
                print(f"Code is {response['error']['code']}", flush=True)
                print(response['error']['message'], flush=True)
                return "rate_limit_exceeded"
                if self.multiple_api_keys == True:
                    new_key = pick_next_item(self.current_key_file, self.key_lists)
                    self.update_key(new_key)
                    self.current_key_file = new_key #override key
                    print("New key is from the file: ", new_key)
            else:
                print("Code is different")
                print(response)
                print(f"{response['error']['code']=}")
        return ""

    def update_key(self, key, load_from_file=True):
        if load_from_file:
            self.api_key = get_api_key(key)
        else:
            self.api_key = key

class GPT4o(GPT4v):
    def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4o-2024-05-13"):
        super().__init__(key, url, are_images_encoded, model_name)

if __name__ == "__main__":
    model = GPT4o('sk-cB6h7HcCSDIp71gs6lFLZxKE0dOYOnJbxzES6kWXe1Wb2VHS', model_name="gpt-4.1")
    prompt = model.prepare_prompt(['https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/DiffEdit/sample_34_1.jpg', 'https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/input/sample_34_1.jpg'], 'What is difference between two images?')
    print("prompt : \n", prompt)
    res = model.get_parsed_output(prompt)
    print("result : \n", res)

	import base64
	import requests
	from io import BytesIO, StringIO
	from typing import Union, Optional, Tuple, List
	from PIL import Image, ImageOps
	import os

	def get_api_key(file_path):
	# Read the API key from the first line of the file
	with open(file_path, 'r') as file:
	return file.readline().strip()

	# Function to encode the image
	def encode_image(image_path):
	with open(image_path, "rb") as image_file:
	return base64.b64encode(image_file.read()).decode('utf-8')

	def pick_next_item(current_item, item_list):
	if current_item not in item_list:
	raise ValueError("Current item is not in the list")
	current_index = item_list.index(current_item)
	next_index = (current_index + 1) % len(item_list)

	return item_list[next_index]

	# Function to encode a PIL image
	def encode_pil_image(pil_image):
	# Create an in-memory binary stream
	image_stream = BytesIO()

	# Save the PIL image to the binary stream in JPEG format (you can change the format if needed)
	pil_image.save(image_stream, format='JPEG')

	# Get the binary data from the stream and encode it as base64
	image_data = image_stream.getvalue()
	base64_image = base64.b64encode(image_data).decode('utf-8')

	return base64_image


	def load_image(image: Union[str, Image.Image], format: str = "RGB", size: Optional[Tuple] = None) -> Image.Image:
	"""
	Load an image from a given path or URL and convert it to a PIL Image.

	Args:
	image (Union[str, Image.Image]): The image path, URL, or a PIL Image object to be loaded.
	format (str, optional): Desired color format of the resulting image. Defaults to "RGB".
	size (Optional[Tuple], optional): Desired size for resizing the image. Defaults to None.

	Returns:
	Image.Image: A PIL Image in the specified format and size.

	Raises:
	ValueError: If the provided image format is not recognized.
	"""
	if isinstance(image, str):
	if image.startswith("http://") or image.startswith("https://"):
	image = Image.open(requests.get(image, stream=True).raw)
	elif os.path.isfile(image):
	image = Image.open(image)
	else:
	raise ValueError(
	f"Incorrect path or url, URLs must start with `http://` or `https://`, and {image} is not a valid path"
	)
	elif isinstance(image, Image.Image):
	image = image
	else:
	raise ValueError(
	"Incorrect format used for image. Should be an url linking to an image, a local path, or a PIL image."
	)
	image = ImageOps.exif_transpose(image)
	image = image.convert(format)
	if (size != None):
	image = image.resize(size, Image.LANCZOS)
	return image

	class GPT4v():
	def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4-vision-preview"):
	"""OpenAI GPT-4-vision model wrapper
	Args:
	api_key_path (str): Path to the API key file. Defaults to 'keys/secret.env'.
	are_images_encoded (bool): Whether the images are encoded in base64. Defaults to False.
	"""
	self.multiple_api_keys = False
	self.current_key_file = None
	self.api_key = key

	self.url = url
	self.model_name = model_name
	self.use_encode = are_images_encoded

	def prepare_input(self, image_links: List = [], text_prompt: str = ""):
	prompt_content = []
	text_dict = {
	"type": "text",
	"text": text_prompt
	}
	prompt_content.append(text_dict)

	if not isinstance(image_links, list):
	image_links = [image_links]

	for image_link in image_links:
	image = load_image(image_link)
	if self.use_encode:
	visual_dict = {
	"type": "image_url",
	"image_url": {"url": f"data:image/jpeg;base64,{encode_pil_image(image)}"}
	}
	else:
	visual_dict = {
	"type": "image_url",
	"image_url": {"url": image_link}
	}
	prompt_content.append(visual_dict)
	return prompt_content

	def inference(self, prompt, seed: Optional[int] = None):
	payload = {
	"model": self.model_name,
	"messages": [
	{
	"role": "user",
	"content": prompt
	}
	],
	# "max_tokens": 1400
	}
	headers = {
	"Content-Type": "application/json",
	"Authorization": f"Bearer {self.api_key}"
	}
	# try:
	response = requests.post(self.url, json=payload, headers=headers, timeout=180) # Set timeout to 5 minutes (300 seconds)
	# except Exception as e:
	# print(f"Error: {e}")
	# return ""
	#return response.text
	return self.extract_response(response)

	def extract_response(self, response):
	try:
	response = response.json()
	out = response['choices'][0]['message']['content']
	return out
	except Exception as e:
	print(f"Error: {e}")
	if response['error']['code'] == 'content_policy_violation':
	print("Code is content_policy_violation")
	elif response['error']['code'] in ['rate_limit_exceeded', 'insufficient_quota', 'insufficient_user_quota']:
	print(f"Code is {response['error']['code']}", flush=True)
	print(response['error']['message'], flush=True)
	return "rate_limit_exceeded"
	if self.multiple_api_keys == True:
	new_key = pick_next_item(self.current_key_file, self.key_lists)
	self.update_key(new_key)
	self.current_key_file = new_key #override key
	print("New key is from the file: ", new_key)
	else:
	print("Code is different")
	print(response)
	print(f"{response['error']['code']=}")
	return ""

	def update_key(self, key, load_from_file=True):
	if load_from_file:
	self.api_key = get_api_key(key)
	else:
	self.api_key = key

	class GPT4o(GPT4v):
	def __init__(self, key, url="https://api.openai.com/v1/chat/completions", are_images_encoded=False, model_name="gpt-4o-2024-05-13"):
	super().__init__(key, url, are_images_encoded, model_name)

	if __name__ == "__main__":
	model = GPT4o('sk-cB6h7HcCSDIp71gs6lFLZxKE0dOYOnJbxzES6kWXe1Wb2VHS', model_name="gpt-4.1")
	prompt = model.prepare_prompt(['https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/DiffEdit/sample_34_1.jpg', 'https://chromaica.github.io/Museum/ImagenHub_Text-Guided_IE/input/sample_34_1.jpg'], 'What is difference between two images?')
	print("prompt : \n", prompt)
	res = model.get_parsed_output(prompt)
	print("result : \n", res)