diff --git a/PixArtAlpha.py b/PixArtAlpha.py new file mode 100644 index 0000000..638ca99 --- /dev/null +++ b/PixArtAlpha.py @@ -0,0 +1,163 @@ +import torch +import os +import folder_paths +from diffusers import PixArtAlphaPipeline, DPMSolverMultistepScheduler +from huggingface_hub import hf_hub_download +import numpy as np +from .sa_solver_diffusers import SASolverScheduler +from .style_template import style_list + + +device = "cuda" if torch.cuda.is_available() else "cpu" + + +# Create a dictionary from style_list for easier access +styles = {style['name']: (style['prompt'], style['negative_prompt']) for style in style_list} + +STYLE_NAMES = [style['name'] for style in style_list] +DEFAULT_STYLE_NAME = "(No style)" + +def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]: + # Get the prompts for the given style_name, or defaults to the first style if not found + p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) + return p.replace("{prompt}", positive), n + ' ' + negative + + +class PA_Styler_Zho: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"default": "A alpaca made of colorful building blocks, cyberpunk", "multiline": True}), + "negative_prompt": ("STRING", {"default": "asymmetry, worst quality, low quality", "multiline": True}), + "style_name": (STYLE_NAMES, {"default": DEFAULT_STYLE_NAME}) + } + } + + RETURN_TYPES = ('STRING', 'STRING',) + RETURN_NAMES = ('positive_prompt', 'negative_prompt',) + FUNCTION = "prompt_style" + CATEGORY = "🖼️PixArtAlpha" + + def prompt_style(self, style_name, prompt, negative_prompt): + prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt) + + return prompt, negative_prompt + + +class PA_BaseModelLoader_fromhub_Zho: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "base_model_path": ("STRING", {"default": "PixArt-alpha/PixArt-XL-2-1024-MS"}) + } + } + + RETURN_TYPES = ("PAMODEL",) + RETURN_NAMES = ("pipe",) + FUNCTION = "load_model" + CATEGORY = "🖼️PixArtAlpha" + + def load_model(self, base_model_path): + pipe = PixArtAlphaPipeline.from_pretrained( + base_model_path, + torch_dtype=torch.float16, + ).to(device) + return [pipe] + + +class PA_Generation_Zho: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "pipe": ("PAMODEL",), + "positive": ("STRING", {"multiline": True, "forceInput": True}), + "negative": ("STRING", {"multiline": True, "forceInput": True}), + "width": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 32}), + "height": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 32}), + "steps": ("INT", {"default": 20, "min": 1, "max": 100, "step": 1}), + "guidance_scale": ("FLOAT", {"default": 4.5, "min": 0, "max": 20}), + "schedule": (["DPM-Solver", "SA-Solver"],), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate_image" + CATEGORY = "🖼️PixArtAlpha" + + def generate_image(self, positive, negative, pipe, steps, guidance_scale, seed, width, height, schedule): + + if schedule == 'DPM-Solver': + if not isinstance(pipe.scheduler, DPMSolverMultistepScheduler): + pipe.scheduler = DPMSolverMultistepScheduler() + elif schedule == "SA-Solver": + if not isinstance(pipe.scheduler, SASolverScheduler): + pipe.scheduler = SASolverScheduler.from_config(pipe.scheduler.config, algorithm_type='data_prediction', tau_func=lambda t: 1 if 200 <= t <= 800 else 0, predictor_order=2, corrector_order=2) + else: + raise ValueError(f"Unknown schedule: {schedule}") + + generator = torch.Generator(device=device).manual_seed(seed) + + output = pipe( + prompt=positive, + negative_prompt=negative, + num_inference_steps=steps, + generator=generator, + guidance_scale=guidance_scale, + width=width, + height=height, + ) + + # 检查输出类型并相应处理 + if isinstance(output, tuple): + # 当返回的是元组时,第一个元素是图像列表 + images_list = output[0] + else: + # 如果返回的是 StableDiffusionXLPipelineOutput,需要从中提取图像 + images_list = output.images + + # 转换图像为 torch.Tensor,并调整维度顺序为 NHWC + images_tensors = [] + for img in images_list: + # 将 PIL.Image 转换为 numpy.ndarray + img_array = np.array(img) + # 转换 numpy.ndarray 为 torch.Tensor + img_tensor = torch.from_numpy(img_array).float() / 255. + # 转换图像格式为 CHW (如果需要) + if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3: + img_tensor = img_tensor.permute(2, 0, 1) + # 添加批次维度并转换为 NHWC + img_tensor = img_tensor.unsqueeze(0).permute(0, 2, 3, 1) + images_tensors.append(img_tensor) + + if len(images_tensors) > 1: + output_image = torch.cat(images_tensors, dim=0) + else: + output_image = images_tensors[0] + + return (output_image,) + + +NODE_CLASS_MAPPINGS = { + "PA_Styler_Zho": PA_Styler_Zho, + "PA_BaseModelLoader_fromhub_Zho": PA_BaseModelLoader_fromhub_Zho, + "PA_Generation_Zho": PA_Generation_Zho +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PA_Styler_Zho": "🖼️PixArtAlpha Styler", + "PA_BaseModelLoader_fromhub_Zho": "🖼️PixArtAlpha ModelLoader", + "PA_Generation_Zho": "🖼️PixArtAlpha Generation" +}