106 lines
3.5 KiB
Python
106 lines
3.5 KiB
Python
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import node_helpers
|
|
from nodes import common_ksampler
|
|
|
|
def phi2tensor(img):
|
|
tensor = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return tensor
|
|
|
|
def tensor_to_image(tensor):
|
|
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
|
image = Image.fromarray(image_np, mode='RGB')
|
|
return image
|
|
class CFGenerator:
|
|
def __init__(self, model,clip,ae,model_type, device: str):
|
|
self.device = torch.device(device)
|
|
self.model=model
|
|
self.ae=ae
|
|
self.clip=clip
|
|
self.model_type=model_type
|
|
|
|
@torch.inference_mode()
|
|
def generate_image(
|
|
self,
|
|
width,
|
|
height,
|
|
num_steps,
|
|
cfg,
|
|
guidance,
|
|
seed,
|
|
prompt,
|
|
negative_prompt,
|
|
cf_scheduler,
|
|
denoise,
|
|
image,
|
|
):
|
|
if isinstance(prompt,str):
|
|
text=prompt
|
|
elif isinstance(prompt,list):
|
|
text=prompt[0] if len(prompt)==1 else str(prompt)
|
|
else:
|
|
raise "prompt must be str or list"
|
|
if isinstance(negative_prompt, str):
|
|
negative_text = negative_prompt
|
|
elif isinstance(negative_prompt, list):
|
|
negative_text =negative_prompt[0] if len(negative_prompt) == 1 else str(negative_prompt)
|
|
else:
|
|
raise "prompt must be str or list"
|
|
|
|
batch_size=1
|
|
|
|
|
|
#cf clip postive
|
|
tokens = self.clip.tokenize(text)
|
|
output = self.clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
|
|
cond = output.pop("cond")
|
|
conditioning=[[cond, output]]
|
|
|
|
#flux GUIDANCE
|
|
if self.model_type == "FLUX":
|
|
postive_c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
|
|
else:
|
|
postive_c=conditioning
|
|
|
|
#cf neg
|
|
tokens_ = self.clip.tokenize(negative_text)
|
|
output_ = self.clip.encode_from_tokens(tokens_, return_pooled=True, return_dict=True)
|
|
cond_ = output_.pop("cond")
|
|
negative_c = [[cond_, output_]]
|
|
|
|
if image:
|
|
if isinstance(image,list):
|
|
image=image[0] if len(image)==1 else None
|
|
if isinstance(image,torch.Tensor):
|
|
pixels=image
|
|
elif isinstance(image,Image.Image):
|
|
pixels=phi2tensor(image)
|
|
elif isinstance(image, np.ndarray):
|
|
init_image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
|
|
pixels = init_image.unsqueeze(0)
|
|
else:
|
|
raise "image input error"
|
|
#print(type(image),image,pixels.shape,type(pixels))
|
|
t = self.ae.encode(pixels[:, :, :, :3])
|
|
else:
|
|
if self.model_type == "FLUX":
|
|
t = torch.ones([batch_size, 16, height // 8, width // 8], device=self.device) * 0.0609
|
|
else:
|
|
t = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
|
|
|
|
latent = {"samples": t}
|
|
|
|
sampler_name= cf_scheduler["name"]
|
|
scheduler=cf_scheduler["scheduler"]
|
|
samples=common_ksampler(self.model, seed, num_steps, cfg, sampler_name, scheduler, postive_c, negative_c, latent, denoise=denoise)
|
|
samples_=samples[0]["samples"]
|
|
img_out=self.ae.decode(samples_)
|
|
#print(img_out.shape,type(img_out))
|
|
img_pil=tensor_to_image(img_out)
|
|
return img_pil
|
|
|
|
|
|
|