132 lines
3.6 KiB
Python
132 lines
3.6 KiB
Python
import torch
|
|
from .utils import (
|
|
ASPECT_RATIOS,
|
|
add_hf_token,
|
|
call_api_with_fallback,
|
|
generation_url_to_tensor,
|
|
resolve_api_key,
|
|
upload_image,
|
|
)
|
|
|
|
EDIT_MODELS = [
|
|
"p-image-edit",
|
|
]
|
|
|
|
EDIT_ASPECT_RATIOS = ["match_input_image"] + ASPECT_RATIOS
|
|
|
|
EDIT_LORA_MODELS = [
|
|
"p-image-edit-lora",
|
|
]
|
|
|
|
|
|
class PrunaImageEdit:
|
|
"""Edit an image with a text instruction using the Pruna API.
|
|
|
|
Accepts a single ComfyUI IMAGE. The image is uploaded to the Pruna files
|
|
API and passed as a URL to the editing model.
|
|
"""
|
|
|
|
CATEGORY = "pruna ai"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "edit"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (EDIT_MODELS, {"default": "p-image-edit"}),
|
|
"aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
def edit(
|
|
self,
|
|
image: torch.Tensor,
|
|
prompt: str,
|
|
model: str,
|
|
aspect_ratio: str,
|
|
api_key: str,
|
|
):
|
|
key = resolve_api_key(api_key)
|
|
image_url = upload_image(image, key)
|
|
payload = {
|
|
"input": {
|
|
"prompt": prompt,
|
|
"images": [image_url],
|
|
"aspect_ratio": aspect_ratio,
|
|
}
|
|
}
|
|
data = call_api_with_fallback(model, payload, key)
|
|
tensor = generation_url_to_tensor(data)
|
|
return (tensor,)
|
|
|
|
|
|
class PrunaImageEditLoRA:
|
|
"""Edit an image with a text instruction and a custom LoRA using the Pruna API.
|
|
|
|
LoRA weights must be a HuggingFace URL in the format:
|
|
huggingface.co/<username>/<repo>[/<filename>]
|
|
|
|
Note: LoRAs used with p-image-edit-lora must be trained with p-image-edit-trainer.
|
|
"""
|
|
|
|
CATEGORY = "pruna ai"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "edit"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (EDIT_LORA_MODELS, {"default": "p-image-edit-lora"}),
|
|
"lora_weights": (
|
|
"STRING",
|
|
{"default": "huggingface.co/username/repo"},
|
|
),
|
|
"lora_scale": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": -1.0, "max": 3.0, "step": 0.05},
|
|
),
|
|
"aspect_ratio": (EDIT_ASPECT_RATIOS, {"default": "match_input_image"}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"hf_api_token": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
def edit(
|
|
self,
|
|
image: torch.Tensor,
|
|
prompt: str,
|
|
model: str,
|
|
lora_weights: str,
|
|
lora_scale: float,
|
|
aspect_ratio: str,
|
|
api_key: str,
|
|
hf_api_token: str = "",
|
|
):
|
|
key = resolve_api_key(api_key)
|
|
image_url = upload_image(image, key)
|
|
payload: dict = {
|
|
"input": {
|
|
"prompt": prompt,
|
|
"images": [image_url],
|
|
"lora_weights": lora_weights,
|
|
"lora_scale": lora_scale,
|
|
"aspect_ratio": aspect_ratio,
|
|
}
|
|
}
|
|
add_hf_token(payload, hf_api_token)
|
|
|
|
data = call_api_with_fallback(model, payload, key)
|
|
tensor = generation_url_to_tensor(data)
|
|
return (tensor,)
|