38 lines
1.0 KiB
Python
38 lines
1.0 KiB
Python
from .utils import (
|
|
ASPECT_RATIOS,
|
|
call_api_with_fallback,
|
|
generation_url_to_tensor,
|
|
resolve_api_key,
|
|
)
|
|
|
|
T2I_MODELS = [
|
|
"p-image",
|
|
]
|
|
|
|
|
|
class PrunaTextToImage:
|
|
"""Generate an image from a text prompt using the Pruna API."""
|
|
|
|
CATEGORY = "pruna ai"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (T2I_MODELS, {"default": "p-image"}),
|
|
"aspect_ratio": (ASPECT_RATIOS, {"default": "1:1"}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
def generate(self, prompt: str, model: str, aspect_ratio: str, api_key: str):
|
|
key = resolve_api_key(api_key)
|
|
payload = {"input": {"prompt": prompt, "aspect_ratio": aspect_ratio}}
|
|
data = call_api_with_fallback(model, payload, key)
|
|
tensor = generation_url_to_tensor(data)
|
|
return (tensor,)
|