pixaris.generation package

Submodules

pixaris.generation.base module

class pixaris.generation.base.ImageGenerator[source]

Bases: object

When implementing a new Image Generator, inherit from this one and implement all the abstract methods.

abstractmethod generate_single_image(args: dict[str, any]) → tuple[Image, str][source]
validate_inputs_and_parameters(inputs: list[dict] = [], parameters: list[dict] = []) → bool[source]

pixaris.generation.comfyui module

class pixaris.generation.comfyui.ComfyGenerator(workflow_apiformat_json: dict, api_host: str = 'localhost:8188')[source]

Bases: ImageGenerator

ComfyGenerator is a class that extends the ImageGenerator class to provide functionality for generating images using a specified workflow and API host. It uses the ComfyUI API to generate images based on the provided workflow and input parameters.

Parameters:

  • workflow_apiformat_json (dict) – The workflow file in API format.

  • api_host (str) – The API host URL. For local experimenting, put “localhost:8188”. There has to be a tunnel to a running comfyUI instance active on port 8188

generate_single_image(args: dict[str, any]) → tuple[Image, str][source]

Generates a single image based on the provided arguments. For this it modifies and executed the workflow to generate the image.

Parameters:

args – A dictionary containing the following keys:

  • “workflow_apiformat_json” (dict): The workflow file in JSON apiformat.

  • “pillow_images” (list[dict]): A dict of [str, Image.Image]. The keys should be Node names The values should be the PIL Image objects to be loaded. Should look like this:

    "pillow_images": [{
"node_name": "Load Input Image",
"pillow_image": Image.new("RGB", (100, 100), color="red")}]

  • “generation_params” (list[dict]): A dictionary of generation_params for the image generation process. Should look like this:

    "generation_params": [{
"node_name": "GroundingDinoSAMSegment (segment anything)",
"input": "prompt",
"value": "model, bag, hair"}]
Rtype args:

dict[str, any]

Returns:

The generated image and its name

Return type:

tuple[Image.Image, str]

validate_inputs_and_parameters(dataset: List[dict[str, List[dict[str, Image]]]] = [], args: dict[str, any] = {}) → str[source]

Validates the workflow file to ensure that it is in the correct format.

Parameters:

  • dataset (List[dict[str, List[dict[str, Image.Image]]]) – A list of dictionaries containing the images to be loaded.

  • args (dict[str, any]) – A dictionary containing the parameters to be used for the image generation process.

Returns:

The path to the validated workflow file.

Return type:

str

pixaris.generation.comfyui_cluster module

class pixaris.generation.comfyui_cluster.ComfyClusterGenerator(workflow_apiformat_json: str)[source]

Bases: ImageGenerator

Cluster to run Comfy workflows. It will automatically fetch available hosts, initiate a new ComfyGenerator for each and distribute the workflows to them. If the environment variable DEV_MODE is set to true, it will run the workflows locally and it uses localhost:8188.

Parameters:

workflow_apiformat_json (str) – The path to the workflow file in API format. (ABSOLUTE PATH)!

close()[source]

Close the cluster and release all hosts.

generate_single_image(args: dict[str, any]) → tuple[Image, str][source]

Generates a single image based on the provided arguments. For this it searches for a host, initialises a ComfyGenerator, and lets it modify and execute the workflow to generate the image.

Parameters:

args – A dictionary containing the following keys:

  • “workflow_apiformat_json” (str): The path to the workflow file in API format. (ABSOLUTE PATH)!

  • “pillow_images” (list[dict]): A dict of [str, Image.Image]. The keys should be Node names The values should be the PIL Image objects to be loaded. Should look like this:

    "pillow_images": [{
"node_name": "Load Input Image",
"pillow_image": Image.new("RGB", (100, 100), color="red"),}]

  • “generation_params” (list[dict]): A dictionary of generation_params for the image generation process. Should look like this:

    "generation_params": [{"node_name": "GroundingDinoSAMSegment (segment anything)",
"input": "prompt",
"value": "model, bag, hair"}],
Rtype args:

dict[str, any]

Returns:

The generated image.

Return type:

tuple[Image.Image, str]

start_background_task()[source]

Start a background task to update the available hosts every minute.

update_available_hosts()[source]

Update the available hosts by fetching the IPs of the pods and checking if the Comfy UI is running on them. Use mutex to avoid conflicts.

validate_inputs_and_parameters(dataset: List[dict[str, List[dict[str, Image]]]] = [], args: dict[str, any] = {}) → str[source]

Creates a dummy generator and then calls ComfyGenerator’s validate_inputs_and_parameters method.

Parameters:

  • dataset (List[dict[str, List[dict[str, Image.Image]]]) – A list of dictionaries containing the images to be loaded.

  • args (dict[str, any]) – A dictionary containing the parameters to be used for the image generation process.

Returns:

The path to the validated workflow file.

Return type:

str

pixaris.generation.flux module

class pixaris.generation.flux.FluxFillGenerator[source]

Bases: ImageGenerator

FluxFillGenerator is responsible for generating images using the Flux API, specifically the fill model, which needs an image and a mask as input.

generate_single_image(args: dict[str, any]) → tuple[Image, str][source]

Generates a single image based on the provided arguments.

Parameters:

args – A dictionary containing the following keys:

  • pillow_images (list[dict]): A list of dictionaries containing pillow images and mask images.

  • generation_params (list[dict]): A list of dictionaries containing generation params.

Returns:

A tuple containing:

  • image (Image.Image): The generated image.

  • image_name (str): The name of the generated image.

Return type:

tuple[Image.Image, str]

validate_inputs_and_parameters(dataset: List[dict[str, List[dict[str, Image]]]] = [], args: dict[str, any] = {})[source]

Validates the provided dataset and parameters for image generation.

Parameters:

  • dataset (List[dict[str, List[dict[str, Image.Image]]]) – A list of datasets containing image and mask information.

  • args (dict[str, any]) – A dictionary containing the parameters to be used for the image generation process.

Raises:

ValueError – If the validation fails for any reason (e.g., missing fields).

pixaris.generation.gemini module

class pixaris.generation.gemini.GeminiGenerator(gcp_project_id: str, gcp_location: str = 'us-central1', model_name: str = 'gemini-2.0-flash-exp', verbose: bool = False)[source]

Bases: ImageGenerator

GeminiGenerator is a class that generates images using the Google Gemini API.

Parameters:

  • gcp_project_id (str) – The Google Cloud Platform project ID.

  • gcp_location (str) – The Google Cloud Platform location.

generate_single_image(args: dict[str, any]) → tuple[Image, str][source]

Generates a single image based on the provided arguments.

Parameters:

args – A dictionary containing the following keys:

  • pillow_images (list[dict]): A list of dictionaries containing pillow images.

  • prompt (str): The prompt that should be used for the generation.

Returns:

A tuple containing:

  • image (Image.Image): The generated image.

  • image_name (str): The name of the generated image.

Return type:

tuple[Image.Image, str]

validate_inputs_and_parameters(dataset: List[dict[str, List[dict[str, Image]]]] = [], args: dict[str, any] = {})[source]

Validates the provided dataset and parameters for image generation.

Parameters:

  • dataset (List[dict[str, List[dict[str, Image.Image]]]) – A list of dicts containing image information.

  • args (dict[str, any]) – A dictionary containing the parameters to be used for the image generation process.

Raises:

ValueError – If the validation fails for any reason (e.g., missing fields).

pixaris.generation.imagen2 module

class pixaris.generation.imagen2.Imagen2Generator(gcp_project_id: str, gcp_location: str)[source]

Bases: ImageGenerator

Imagen2Generator is a class that generates images using the Google Imagen API.

Parameters:

  • gcp_project_id (str) – The Google Cloud Platform project ID.

  • gcp_location (str) – The Google Cloud Platform location.

generate_single_image(args: dict[str, any]) → tuple[Image, str][source]

Generates a single image based on the provided arguments.

Parameters:

args – A dictionary containing the following keys:

  • pillow_images (List[dict[str, List[dict[str, Image.Image]]]]): A list of dictionaries containing

    pillow images and mask images.

  • prompt (str): The prompt that should be used for the generation.

Returns:

A tuple containing:

  • image (Image.Image): The generated image.

  • image_name (str): The name of the generated image.

Return type:

tuple[Image.Image, str]

validate_inputs_and_parameters(dataset: List[dict[str, List[dict[str, Image]]]], args: dict[str, any] = {})[source]

Validates the provided dataset and parameters for image generation.

Parameters:

  • dataset (List[dict[str, List[dict[str, Image.Image]]]) – A list of datasets containing image and mask information.

  • args (dict[str, any]) – A dictionary containing the parameters to be used for the image generation process.

Raises:

ValueError – If the validation fails for any reason (e.g., missing fields).

pixaris.generation.utils module

pixaris.generation.utils.encode_image_to_bytes(pillow_image: Image) → bytes[source]

Encodes a PIL image to bytes.

Parameters:

pillow_image (PIL.Image.Image) – The PIL image.

Returns:

Byte array representation of the image.

Return type:

bytes

pixaris.generation.utils.extract_value_from_list_of_dicts(dict_list, identifying_key, identifying_value, return_key, default_value=None)[source]

Extracts a value from a list of dictionaries based on a key-value pair. This function searches through a list of dictionaries and returns the value associated with a specified key for the first dictionary that matches a given key-value pair. If no such dictionary is found, it returns a default value. If the default value is not provided and no matching dictionary is found, a ValueError is raised.

Parameters:

  • dict_list (list[dict]) – A list of dictionaries to search through.

  • identifying_key (_any) – The key to identify the dictionary.

  • identifying_value (_any) – The value to match against the identifying key.

  • return_key (_any) – The key whose value is to be returned from the matching dictionary.

  • default_value (_any, optional) – The value to return if no matching dictionary is found. If not provided, defaults to None.

Raises:

ValueError – If no matching dictionary is found and no default value is provided.

Returns:

The value associated with the return key from the matching dictionary, or the default value if no match is found.

Return type:

_any

Module contents