File size: 2,885 Bytes
8247a04 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from huggingface_hub import InferenceClient
from PIL import Image
import io
import config
class DiffusionInference:
def __init__(self, api_key=None):
"""
Initialize the inference client with the Hugging Face API token.
"""
self.api_key = api_key or config.HF_TOKEN
self.client = InferenceClient(
provider="hf-inference",
api_key=self.api_key,
)
def text_to_image(self, prompt, model_name=None, negative_prompt=None, **kwargs):
"""
Generate an image from a text prompt.
Args:
prompt (str): The text prompt to guide image generation
model_name (str, optional): The model to use for inference
negative_prompt (str, optional): What not to include in the image
**kwargs: Additional parameters to pass to the model
Returns:
PIL.Image: The generated image
"""
model = model_name or config.DEFAULT_TEXT2IMG_MODEL
# Set up parameters dictionary
params = {"prompt": prompt}
if negative_prompt:
params["negative_prompt"] = negative_prompt
# Add any additional parameters
params.update(kwargs)
try:
image = self.client.text_to_image(model=model, **params)
return image
except Exception as e:
print(f"Error generating image: {e}")
raise
def image_to_image(self, image, prompt=None, model_name=None, negative_prompt=None, **kwargs):
"""
Generate a new image from an input image and optional prompt.
Args:
image (PIL.Image or str): Input image or path to image
prompt (str, optional): Text prompt to guide the transformation
model_name (str, optional): The model to use for inference
negative_prompt (str, optional): What not to include in the image
**kwargs: Additional parameters to pass to the model
Returns:
PIL.Image: The generated image
"""
model = model_name or config.DEFAULT_IMG2IMG_MODEL
# Convert image path to PIL Image if needed
if isinstance(image, str):
image = Image.open(image)
# Set up parameters dictionary
params = {"image": image}
if prompt:
params["prompt"] = prompt
if negative_prompt:
params["negative_prompt"] = negative_prompt
# Add any additional parameters
params.update(kwargs)
try:
result = self.client.image_to_image(model=model, **params)
return result
except Exception as e:
print(f"Error transforming image: {e}")
raise
|