Spaces:
Running
Running
| import json | |
| import io | |
| import random | |
| import gradio as gr | |
| from PIL import Image | |
| from generate import * | |
| from typing import Dict, Any | |
| from processImage import process_and_encode_image | |
| def display_image(image_bytes): | |
| if isinstance(image_bytes, str): | |
| # If we received a string (error message), return it to be displayed | |
| return None, gr.update(visible=True, value=image_bytes) | |
| elif image_bytes: | |
| # If we received image bytes, process and display the image | |
| return Image.open(io.BytesIO(image_bytes)), gr.update(visible=False) | |
| else: | |
| # Handle None case | |
| return None, gr.update(visible=False) | |
| def process_optional_params(**kwargs) -> Dict[str, Any]: | |
| return {k: v for k, v in kwargs.items() if v is not None} | |
| def process_images(primary=None, secondary=None, validate=True) -> Dict[str, str]: | |
| if validate and primary is None: | |
| raise ValueError("Primary image is required.") | |
| result = {} | |
| if primary: | |
| result["image"] = process_and_encode_image(primary) | |
| if secondary: | |
| result["maskImage"] = process_and_encode_image(secondary) | |
| return result | |
| def create_image_generation_config(height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| return { | |
| "numberOfImages": 1, | |
| "height": height, | |
| "width": width, | |
| "quality": quality, | |
| "cfgScale": cfg_scale, | |
| "seed": seed | |
| } | |
| def build_request(task_type, params, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| param_dict = {"TEXT_IMAGE": "textToImageParams", "INPAINTING": "inPaintingParams", | |
| "OUTPAINTING":"outPaintingParams","IMAGE_VARIATION":"imageVariationParams", | |
| "COLOR_GUIDED_GENERATION":"colorGuidedGenerationParams","BACKGROUND_REMOVAL":"backgroundRemovalParams"} | |
| return json.dumps({ | |
| "taskType": task_type, | |
| param_dict[task_type]: params, | |
| "imageGenerationConfig": create_image_generation_config( | |
| height=height, | |
| width=width, | |
| quality=quality, | |
| cfg_scale=cfg_scale, | |
| seed=seed | |
| ) | |
| }) | |
| def check_return(result): | |
| if not isinstance(result, bytes): | |
| return None, gr.update(visible=True, value=result) | |
| return Image.open(io.BytesIO(result)), gr.update(visible=False) | |
| def text_to_image(prompt, negative_text=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| text_to_image_params = {"text": prompt, | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| } | |
| body = build_request("TEXT_IMAGE", text_to_image_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def inpainting(image, mask_prompt=None, mask_image=None, text=None, negative_text=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| images = process_images(primary=image, secondary=None) | |
| for value in images.values(): | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| # Prepare the inPaintingParams dictionary | |
| if mask_prompt and mask_image: | |
| raise ValueError("You must specify either maskPrompt or maskImage, but not both.") | |
| if not mask_prompt and not mask_image: | |
| raise ValueError("You must specify either maskPrompt or maskImage.") | |
| # Prepare the inPaintingParams dictionary with the appropriate mask parameter | |
| in_painting_params = { | |
| **images, # Unpacks image and maskImage if present | |
| **({"maskPrompt": mask_prompt} if mask_prompt not in [None, ""] else {}), | |
| **({"text": text} if text not in [None, ""] else {}), | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| } | |
| body = build_request("INPAINTING", in_painting_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def outpainting(image, mask_prompt=None, mask_image=None, text=None, negative_text=None, outpainting_mode="DEFAULT", height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| images = process_images(primary=image, secondary=None) | |
| for value in images.values(): | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| if mask_prompt and mask_image: | |
| raise ValueError("You must specify either maskPrompt or maskImage, but not both.") | |
| if not mask_prompt and not mask_image: | |
| raise ValueError("You must specify either maskPrompt or maskImage.") | |
| # Prepare the outPaintingParams dictionary | |
| out_painting_params = { | |
| **images, # Unpacks image and maskImage if present | |
| **process_optional_params( | |
| **({"maskPrompt": mask_prompt} if mask_prompt not in [None, ""] else {}), | |
| **({"text": text} if text not in [None, ""] else {}), | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| ) | |
| } | |
| body = build_request("OUTPAINTING", out_painting_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def image_variation(images, text=None, negative_text=None, similarity_strength=0.5, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| encoded_images = [] | |
| for image_path in images: | |
| with open(image_path, "rb") as image_file: | |
| value = process_and_encode_image(image_file) | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| encoded_images.append(value) | |
| # Prepare the imageVariationParams dictionary | |
| image_variation_params = { | |
| "images": encoded_images, | |
| **({"text": text} if text not in [None, ""] else {}), | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| } | |
| body = build_request("IMAGE_VARIATION", image_variation_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def image_conditioning(condition_image, text, negative_text=None, control_mode="CANNY_EDGE", control_strength=0.7, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| condition_image_encoded = process_images(primary=condition_image) | |
| for value in condition_image_encoded.values(): | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| # Prepare the textToImageParams dictionary | |
| text_to_image_params = { | |
| "text": text, | |
| "controlMode": control_mode, | |
| "controlStrength": control_strength, | |
| "conditionImage": condition_image_encoded.get('image'), | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| } | |
| body = build_request("TEXT_IMAGE", text_to_image_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def color_guided_content(text=None, reference_image=None, negative_text=None, colors=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0): | |
| reference_image_str = None | |
| if reference_image is not None and not isinstance(reference_image, type(None)): | |
| reference_image_encoded = process_images(primary=reference_image) | |
| for value in reference_image_encoded.values(): | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| reference_image_str = reference_image_encoded.get('image') | |
| if not colors: | |
| colors = "#FF5733,#33FF57,#3357FF,#FF33A1,#33FFF5,#FF8C33,#8C33FF,#33FF8C,#FF3333,#33A1FF" | |
| color_guided_generation_params = { | |
| "text": text, | |
| "colors": colors.split(','), | |
| **({"referenceImage": reference_image_str} if reference_image_str is not None else {}), | |
| **({"negativeText": negative_text} if negative_text not in [None, ""] else {}) | |
| } | |
| body = build_request("COLOR_GUIDED_GENERATION", color_guided_generation_params, height, width, quality, cfg_scale, seed) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def background_removal(image): | |
| input_image = process_images(primary=image) | |
| for value in input_image.values(): | |
| if len(value) < 200: | |
| return None, gr.update(visible=True, value=value) | |
| body = json.dumps({ | |
| "taskType": "BACKGROUND_REMOVAL", | |
| "backgroundRemovalParams": { | |
| "image": input_image.get('image') | |
| } | |
| }) | |
| result = generate_image(body) | |
| return check_return(result) | |
| def generate_nova_prompt(): | |
| with open('seeds.json', 'r') as file: | |
| data = json.load(file) | |
| if 'seeds' not in data or not isinstance(data['seeds'], list): | |
| raise ValueError("The JSON file must contain a 'seeds' key with a list of strings.") | |
| random_string = random.choice(data['seeds']) | |
| prompt = f""" | |
| Generate a creative image prompt that builds upon this concept: "{random_string}" | |
| Requirements: | |
| - Create a new, expanded prompt without mentioning or repeating the original concept | |
| - Focus on vivid visual details and artistic elements | |
| - Keep the prompt under 1000 characters | |
| - Do not include any meta-instructions or seed references | |
| - Return only the new prompt text | |
| Response Format: | |
| [Just the new prompt text, nothing else] | |
| """ | |
| messages = [ | |
| {"role": "user", "content": [{"text": prompt}]} | |
| ] | |
| return generate_prompt(messages) | |