diffused.generate

 1from typing import NotRequired, TypedDict, Unpack
 2
 3import diffusers
 4import torch
 5from PIL import Image
 6
 7
 8class Generate(TypedDict):
 9    model: str
10    prompt: str
11    negative_prompt: NotRequired[str]
12    image: NotRequired[str]
13    mask_image: NotRequired[str]
14    width: NotRequired[int]
15    height: NotRequired[int]
16    num_images_per_prompt: NotRequired[int]
17    guidance_scale: NotRequired[float]
18    num_inference_steps: NotRequired[int]
19    strength: NotRequired[float]
20    seed: NotRequired[int]
21    device: NotRequired[str]
22    use_safetensors: NotRequired[bool]
23
24
25def generate(**kwargs: Unpack[Generate]) -> list[Image.Image]:
26    """
27    Generate image with diffusion model.
28
29    Args:
30        model (str): Diffusion model.
31        prompt (str): Text prompt.
32        negative_prompt (str): What to exclude from the generated image.
33        image (str): Input image path or URL.
34        mask_image (str): Mask image path or URL.
35        width (int): Generated image width in pixels.
36        height (int): Generated image height in pixels.
37        num_images_per_prompt (int): Number of images per prompt.
38        guidance_scale (float): How much the prompt influences image generation.
39        num_inference_steps (int): Number of diffusion steps used for generation.
40        strength (float): How much noise is added to the input image.
41        seed (int): Seed for generating reproducible images.
42        device (str): Device to accelerate computation (cpu, cuda, mps).
43        use_safetensors (bool): Whether to load safetensors.
44
45    Returns:
46        images (list[PIL.Image.Image]): Pillow images.
47    """
48    pipeline_args = {
49        "prompt": kwargs.get("prompt"),
50        "negative_prompt": kwargs.get("negative_prompt"),
51        "width": kwargs.get("width"),
52        "height": kwargs.get("height"),
53        "num_images_per_prompt": kwargs.get("num_images_per_prompt", 1),
54        "use_safetensors": kwargs.get("use_safetensors", True),
55    }
56
57    guidance_scale = kwargs.get("guidance_scale")
58    if guidance_scale is not None and guidance_scale >= 0:
59        pipeline_args["guidance_scale"] = guidance_scale
60
61    num_inference_steps = kwargs.get("num_inference_steps")
62    if num_inference_steps is not None and num_inference_steps >= 0:
63        pipeline_args["num_inference_steps"] = num_inference_steps
64
65    strength = kwargs.get("strength")
66    if strength is not None and strength >= 0:
67        pipeline_args["strength"] = strength
68
69    Pipeline = diffusers.AutoPipelineForText2Image
70
71    if kwargs.get("image"):
72        pipeline_args["image"] = diffusers.utils.load_image(kwargs.get("image"))
73        Pipeline = diffusers.AutoPipelineForImage2Image
74
75    if kwargs.get("mask_image"):
76        pipeline_args["mask_image"] = diffusers.utils.load_image(
77            kwargs.get("mask_image")
78        )
79        Pipeline = diffusers.AutoPipelineForInpainting
80
81    pipeline = Pipeline.from_pretrained(kwargs.get("model"))
82
83    device = kwargs.get("device")
84    if device:
85        pipeline.to(device)
86
87    seed = kwargs.get("seed")
88    if isinstance(seed, int):
89        pipeline_args["generator"] = torch.Generator(device=device).manual_seed(seed)
90
91    images = pipeline(**pipeline_args).images
92    return images