129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
import torch
|
|
from typing import Optional, Union, List, Tuple
|
|
from diffusers.pipelines import FluxPipeline
|
|
from PIL import Image, ImageFilter
|
|
import numpy as np
|
|
import cv2
|
|
|
|
condition_dict = {
|
|
"depth": 0,
|
|
"canny": 1,
|
|
"subject": 4,
|
|
"coloring": 6,
|
|
"deblurring": 7,
|
|
"fill": 9,
|
|
}
|
|
|
|
|
|
class Condition(object):
|
|
def __init__(
|
|
self,
|
|
condition_type: str,
|
|
raw_img: Union[Image.Image, torch.Tensor] = None,
|
|
condition: Union[Image.Image, torch.Tensor] = None,
|
|
mask=None,
|
|
) -> None:
|
|
self.condition_type = condition_type
|
|
assert raw_img is not None or condition is not None
|
|
if raw_img is not None:
|
|
self.condition = self.get_condition(condition_type, raw_img)
|
|
else:
|
|
self.condition = condition
|
|
# TODO: Add mask support
|
|
assert mask is None, "Mask not supported yet"
|
|
|
|
def get_condition(
|
|
self, condition_type: str, raw_img: Union[Image.Image, torch.Tensor]
|
|
) -> Union[Image.Image, torch.Tensor]:
|
|
"""
|
|
Returns the condition image.
|
|
"""
|
|
if condition_type == "depth":
|
|
from transformers import pipeline
|
|
import folder_paths
|
|
import os
|
|
|
|
depth_pipe = pipeline(
|
|
task="depth-estimation",
|
|
# model="LiheYoung/depth-anything-small-hf",
|
|
model=os.path.join(folder_paths.models_dir, 'flux', 'OminiControl', 'depth-anything-small-hf'),
|
|
device="cuda",
|
|
)
|
|
source_image = raw_img.convert("RGB")
|
|
condition_img = depth_pipe(source_image)["depth"].convert("RGB")
|
|
return condition_img
|
|
elif condition_type == "canny":
|
|
img = np.array(raw_img)
|
|
edges = cv2.Canny(img, 100, 200)
|
|
edges = Image.fromarray(edges).convert("RGB")
|
|
return edges
|
|
elif condition_type == "subject":
|
|
return raw_img
|
|
elif condition_type == "coloring":
|
|
return raw_img.convert("L").convert("RGB")
|
|
elif condition_type == "deblurring":
|
|
condition_image = (
|
|
raw_img.convert("RGB")
|
|
.filter(ImageFilter.GaussianBlur(10))
|
|
.convert("RGB")
|
|
)
|
|
return condition_image
|
|
elif condition_type == "fill":
|
|
return raw_img.convert("RGB")
|
|
return self.condition
|
|
|
|
@property
|
|
def type_id(self) -> int:
|
|
"""
|
|
Returns the type id of the condition.
|
|
"""
|
|
return condition_dict[self.condition_type]
|
|
|
|
@classmethod
|
|
def get_type_id(cls, condition_type: str) -> int:
|
|
"""
|
|
Returns the type id of the condition.
|
|
"""
|
|
return condition_dict[condition_type]
|
|
|
|
def _encode_image(self, pipe: FluxPipeline, cond_img: Image.Image) -> torch.Tensor:
|
|
"""
|
|
Encodes an image condition into tokens using the pipeline.
|
|
"""
|
|
cond_img = pipe.image_processor.preprocess(cond_img)
|
|
# cond_img = cond_img.to(pipe.device).to(pipe.dtype)
|
|
cond_img = cond_img.to('cuda').to(pipe.dtype)
|
|
cond_img = pipe.vae.encode(cond_img).latent_dist.sample()
|
|
cond_img = (
|
|
cond_img - pipe.vae.config.shift_factor
|
|
) * pipe.vae.config.scaling_factor
|
|
cond_tokens = pipe._pack_latents(cond_img, *cond_img.shape)
|
|
cond_ids = pipe._prepare_latent_image_ids(
|
|
cond_img.shape[0],
|
|
cond_img.shape[2],
|
|
cond_img.shape[3],
|
|
pipe.device,
|
|
pipe.dtype,
|
|
)
|
|
return cond_tokens, cond_ids
|
|
|
|
def encode(self, pipe: FluxPipeline) -> Tuple[torch.Tensor, torch.Tensor, int]:
|
|
"""
|
|
Encodes the condition into tokens, ids and type_id.
|
|
"""
|
|
if self.condition_type in [
|
|
"depth",
|
|
"canny",
|
|
"subject",
|
|
"coloring",
|
|
"deblurring",
|
|
"fill",
|
|
]:
|
|
tokens, ids = self._encode_image(pipe, self.condition)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Condition type {self.condition_type} not implemented"
|
|
)
|
|
type_id = torch.ones_like(ids[:, :1]) * self.type_id
|
|
return tokens, ids, type_id
|