init
This commit is contained in:
smthem
+201
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,52 @@
|
||||
# DreamO
|
||||
|
||||
Official implementation of **[DreamO: A Unified Framework for Image Customization](https://arxiv.org/abs/2504.16915)**
|
||||
|
||||
[](https://arxiv.org/abs/2504.16915) [](https://huggingface.co/spaces/ByteDance/DreamO) <br>
|
||||
|
||||
|
||||
|
||||
|
||||
### :triangular_flag_on_post: Updates
|
||||
* **2025.05.08**: release codes and models
|
||||
* **2025.04.24**: release DreamO tech report.
|
||||
|
||||
## :wrench: Dependencies and Installation
|
||||
```bash
|
||||
# clone DreamO repo
|
||||
git clone https://github.com/bytedance/DreamO.git
|
||||
cd DreamO
|
||||
# create conda env
|
||||
conda create --name dreamo python=3.10
|
||||
# activate env
|
||||
conda activate dreamo
|
||||
# install dependent packages
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
|
||||
## :zap: Quick Inference
|
||||
### Local Gradio Demo
|
||||
```bash
|
||||
python app.py
|
||||
```
|
||||
|
||||
### Online HuggingFace Demo
|
||||
You can try DreamO demo on [HuggingFace](https://huggingface.co/spaces/ByteDance/DreamO).
|
||||
|
||||
|
||||
## Disclaimer
|
||||
|
||||
This project strives to impact the domain of AI-driven image generation positively. Users are granted the freedom to
|
||||
create images using this tool, but they are expected to comply with local laws and utilize it responsibly.
|
||||
The developers do not assume any responsibility for potential misuse by users.
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If DreamO is helpful, please help to ⭐ the repo.
|
||||
|
||||
If you find this project useful for your research, please consider citing our [paper](https://arxiv.org/abs/2504.16915).
|
||||
|
||||
## :e-mail: Contact
|
||||
If you have any comments or questions, please [open a new issue](https://github.com/xxx/xxx/issues/new/choose) or contact [Yanze Wu](https://tothebeginning.github.io/) and [Chong Mou](mailto:eechongm@gmail.com).
|
||||
+361
@@ -0,0 +1,361 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# import argparse
|
||||
|
||||
import cv2
|
||||
# import gradio as gr
|
||||
import numpy as np
|
||||
import torch
|
||||
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
|
||||
from huggingface_hub import hf_hub_download
|
||||
from PIL import Image
|
||||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from .dreamo.dreamo_pipeline import DreamOPipeline
|
||||
from .dreamo.utils import img2tensor, resize_numpy_image_area, tensor2img
|
||||
from .tools import BEN2
|
||||
|
||||
# parser = argparse.ArgumentParser()
|
||||
# parser.add_argument('--port', type=int, default=8080)
|
||||
# args = parser.parse_args()
|
||||
|
||||
|
||||
class Generator:
|
||||
def __init__(self):
|
||||
device = torch.device('cuda')
|
||||
# preprocessing models
|
||||
# background remove model: BEN2
|
||||
self.bg_rm_model = BEN2.BEN_Base().to(device).eval()
|
||||
hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='models')
|
||||
self.bg_rm_model.loadcheckpoints('models/BEN2_Base.pth')
|
||||
# face crop and align tool: facexlib
|
||||
self.face_helper = FaceRestoreHelper(
|
||||
upscale_factor=1,
|
||||
face_size=512,
|
||||
crop_ratio=(1, 1),
|
||||
det_model='retinaface_resnet50',
|
||||
save_ext='png',
|
||||
device=device,
|
||||
)
|
||||
|
||||
# load dreamo
|
||||
model_root = 'black-forest-labs/FLUX.1-dev'
|
||||
dreamo_pipeline = DreamOPipeline.from_pretrained(model_root, torch_dtype=torch.bfloat16)
|
||||
dreamo_pipeline.load_dreamo_model(device, use_turbo=True)
|
||||
self.dreamo_pipeline = dreamo_pipeline.to(device)
|
||||
|
||||
@torch.no_grad()
|
||||
def get_align_face(self, img):
|
||||
# the face preprocessing code is same as PuLID
|
||||
self.face_helper.clean_all()
|
||||
image_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||||
self.face_helper.read_image(image_bgr)
|
||||
self.face_helper.get_face_landmarks_5(only_center_face=True)
|
||||
self.face_helper.align_warp_face()
|
||||
if len(self.face_helper.cropped_faces) == 0:
|
||||
return None
|
||||
align_face = self.face_helper.cropped_faces[0]
|
||||
|
||||
input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
|
||||
input = input.to(torch.device("cuda"))
|
||||
parsing_out = self.face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
|
||||
parsing_out = parsing_out.argmax(dim=1, keepdim=True)
|
||||
bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
|
||||
bg = sum(parsing_out == i for i in bg_label).bool()
|
||||
white_image = torch.ones_like(input)
|
||||
# only keep the face features
|
||||
face_features_image = torch.where(bg, white_image, input)
|
||||
face_features_image = tensor2img(face_features_image, rgb2bgr=False)
|
||||
|
||||
return face_features_image
|
||||
|
||||
|
||||
# generator = Generator()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def dreamo_generate_image(
|
||||
dreamo_pipeline,ref_conds,
|
||||
# ref_image1,
|
||||
# ref_image2,
|
||||
# ref_task1,
|
||||
# ref_task2,
|
||||
prompt,
|
||||
width,
|
||||
height,
|
||||
num_steps,
|
||||
guidance,
|
||||
seed,
|
||||
true_cfg,
|
||||
cfg_start_step,
|
||||
cfg_end_step,
|
||||
neg_prompt,
|
||||
neg_guidance,
|
||||
first_step_guidance,
|
||||
):
|
||||
#print(prompt)
|
||||
# ref_conds = []
|
||||
# #debug_images = []
|
||||
|
||||
# ref_images = [ref_image1, ref_image2]
|
||||
# ref_tasks = [ref_task1, ref_task2]
|
||||
|
||||
# for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
|
||||
# if ref_image is not None:
|
||||
# if ref_task == "id":
|
||||
# ref_image = generator.get_align_face(ref_image)
|
||||
# elif ref_task != "style":
|
||||
# ref_image = generator.bg_rm_model.inference(Image.fromarray(ref_image))
|
||||
# ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
|
||||
# #debug_images.append(ref_image)
|
||||
# ref_image = img2tensor(ref_image, bgr2rgb=False).unsqueeze(0) / 255.0
|
||||
# ref_image = 2 * ref_image - 1.0
|
||||
# ref_conds.append(
|
||||
# {
|
||||
# 'img': ref_image,
|
||||
# 'task': ref_task,
|
||||
# 'idx': idx + 1,
|
||||
# }
|
||||
# )
|
||||
|
||||
seed = int(seed)
|
||||
if seed == -1:
|
||||
seed = torch.Generator(device="cpu").seed()
|
||||
|
||||
image = dreamo_pipeline(
|
||||
prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=num_steps,
|
||||
guidance_scale=guidance,
|
||||
ref_conds=ref_conds,
|
||||
generator=torch.Generator(device="cpu").manual_seed(seed),
|
||||
true_cfg_scale=true_cfg,
|
||||
true_cfg_start_step=cfg_start_step,
|
||||
true_cfg_end_step=cfg_end_step,
|
||||
negative_prompt=neg_prompt,
|
||||
neg_guidance_scale=neg_guidance,
|
||||
first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else guidance,
|
||||
).images[0]
|
||||
|
||||
return image
|
||||
|
||||
|
||||
# _HEADER_ = '''
|
||||
# <div style="text-align: center; max-width: 650px; margin: 0 auto;">
|
||||
# <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">DreamO</h1>
|
||||
# <p style="font-size: 1rem; margin-bottom: 1.5rem;">Paper: <a href='https://arxiv.org/abs/2504.16915' target='_blank'>DreamO: A Unified Framework for Image Customization</a> | Codes: <a href='https://github.com/bytedance/DreamO' target='_blank'>GitHub</a></p>
|
||||
# </div>
|
||||
|
||||
# ❗️❗️❗️**User Guide:**
|
||||
# - The most important thing to do first is to try the examples provided below the demo, which will help you better understand the capabilities of the DreamO model and the types of tasks it currently supports
|
||||
# - For each input, please select the appropriate task type. For general objects, characters, or clothing, choose IP — we will remove the background from the input image. If you select ID, we will extract the face region from the input image (similar to PuLID). If you select Style, the background will be preserved, and you must prepend the prompt with the instruction: 'generate a same style image.' to activate the style task.
|
||||
# - To accelerate inference, we adopt FLUX-turbo LoRA, which reduces the sampling steps from 25 to 12 compared to FLUX-dev. Additionally, we distill a CFG LoRA, achieving nearly a twofold reduction in steps by eliminating the need for true CFG
|
||||
|
||||
# ''' # noqa E501
|
||||
|
||||
# _CITE_ = r"""
|
||||
# If DreamO is helpful, please help to ⭐ the <a href='https://github.com/bytedance/DreamO' target='_blank'> Github Repo</a>. Thanks!
|
||||
# ---
|
||||
|
||||
# 📧 **Contact**
|
||||
# If you have any questions or feedbacks, feel free to open a discussion or contact <b>wuyanze123@gmail.com</b> and <b>eechongm@gmail.com</b>
|
||||
# """ # noqa E501
|
||||
|
||||
|
||||
# def create_demo():
|
||||
|
||||
# with gr.Blocks() as demo:
|
||||
# gr.Markdown(_HEADER_)
|
||||
|
||||
# with gr.Row():
|
||||
# with gr.Column():
|
||||
# with gr.Row():
|
||||
# ref_image1 = gr.Image(label="ref image 1", type="numpy", height=256)
|
||||
# ref_image2 = gr.Image(label="ref image 2", type="numpy", height=256)
|
||||
# with gr.Row():
|
||||
# ref_task1 = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="task for ref image 1")
|
||||
# ref_task2 = gr.Dropdown(choices=["ip", "id", "style"], value="ip", label="task for ref image 2")
|
||||
# prompt = gr.Textbox(label="Prompt", value="a person playing guitar in the street")
|
||||
# width = gr.Slider(768, 1024, 1024, step=16, label="Width")
|
||||
# height = gr.Slider(768, 1024, 1024, step=16, label="Height")
|
||||
# num_steps = gr.Slider(8, 30, 12, step=1, label="Number of steps")
|
||||
# guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Guidance")
|
||||
# seed = gr.Textbox(label="Seed (-1 for random)", value="-1")
|
||||
# with gr.Accordion("Advanced Options", open=False, visible=False):
|
||||
# ref_res = gr.Slider(512, 1024, 512, step=16, label="resolution for ref image")
|
||||
# neg_prompt = gr.Textbox(label="Neg Prompt", value="")
|
||||
# neg_guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Neg Guidance")
|
||||
# true_cfg = gr.Slider(1, 5, 1, step=0.1, label="true cfg")
|
||||
# cfg_start_step = gr.Slider(0, 30, 0, step=1, label="cfg start step")
|
||||
# cfg_end_step = gr.Slider(0, 30, 0, step=1, label="cfg end step")
|
||||
# first_step_guidance = gr.Slider(0, 10, 0, step=0.1, label="first step guidance")
|
||||
# generate_btn = gr.Button("Generate")
|
||||
# gr.Markdown(_CITE_)
|
||||
|
||||
# with gr.Column():
|
||||
# output_image = gr.Image(label="Generated Image", format='png')
|
||||
# debug_image = gr.Gallery(
|
||||
# label="Preprocessing output (including possible face crop and background remove)",
|
||||
# elem_id="gallery",
|
||||
# )
|
||||
# seed_output = gr.Textbox(label="Used Seed")
|
||||
|
||||
# with gr.Row(), gr.Column():
|
||||
# gr.Markdown("## Examples")
|
||||
# example_inps = [
|
||||
# [
|
||||
# 'example_inputs/woman1.png',
|
||||
# 'ip',
|
||||
# 'profile shot dark photo of a 25-year-old female with smoke escaping from her mouth, the backlit smoke gives the image an ephemeral quality, natural face, natural eyebrows, natural skin texture, award winning photo, highly detailed face, atmospheric lighting, film grain, monochrome', # noqa E501
|
||||
# 9180879731249039735,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/man1.png',
|
||||
# 'ip',
|
||||
# 'a man sitting on the cloud, playing guitar',
|
||||
# 1206523688721442817,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/toy1.png',
|
||||
# 'ip',
|
||||
# 'a purple toy holding a sign saying "DreamO", on the mountain',
|
||||
# 1563188099017016129,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/perfume.png',
|
||||
# 'ip',
|
||||
# 'a perfume under spotlight',
|
||||
# 116150031980664704,
|
||||
# ],
|
||||
# ]
|
||||
# gr.Examples(examples=example_inps, inputs=[ref_image1, ref_task1, prompt, seed], label='IP task')
|
||||
|
||||
# example_inps = [
|
||||
# [
|
||||
# 'example_inputs/hinton.jpeg',
|
||||
# 'id',
|
||||
# 'portrait, Chibi',
|
||||
# 5443415087540486371,
|
||||
# ],
|
||||
# ]
|
||||
# gr.Examples(
|
||||
# examples=example_inps,
|
||||
# inputs=[ref_image1, ref_task1, prompt, seed],
|
||||
# label='ID task (similar to PuLID, will only refer to the face)',
|
||||
# )
|
||||
|
||||
# example_inps = [
|
||||
# [
|
||||
# 'example_inputs/mickey.png',
|
||||
# 'style',
|
||||
# 'generate a same style image. A rooster wearing overalls.',
|
||||
# 6245580464677124951,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/mountain.png',
|
||||
# 'style',
|
||||
# 'generate a same style image. A pavilion by the river, and the distant mountains are endless',
|
||||
# 5248066378927500767,
|
||||
# ],
|
||||
# ]
|
||||
# gr.Examples(examples=example_inps, inputs=[ref_image1, ref_task1, prompt, seed], label='Style task')
|
||||
|
||||
# example_inps = [
|
||||
# [
|
||||
# 'example_inputs/shirt.png',
|
||||
# 'example_inputs/skirt.jpeg',
|
||||
# 'ip',
|
||||
# 'ip',
|
||||
# 'A girl is wearing a short-sleeved shirt and a short skirt on the beach.',
|
||||
# 9514069256241143615,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/woman2.png',
|
||||
# 'example_inputs/dress.png',
|
||||
# 'id',
|
||||
# 'ip',
|
||||
# 'the woman wearing a dress, In the banquet hall',
|
||||
# 7698454872441022867,
|
||||
# ],
|
||||
# ]
|
||||
# gr.Examples(
|
||||
# examples=example_inps,
|
||||
# inputs=[ref_image1, ref_image2, ref_task1, ref_task2, prompt, seed],
|
||||
# label='Try-On task',
|
||||
# )
|
||||
|
||||
# example_inps = [
|
||||
# [
|
||||
# 'example_inputs/dog1.png',
|
||||
# 'example_inputs/dog2.png',
|
||||
# 'ip',
|
||||
# 'ip',
|
||||
# 'two dogs in the jungle',
|
||||
# 6187006025405083344,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/woman3.png',
|
||||
# 'example_inputs/cat.png',
|
||||
# 'ip',
|
||||
# 'ip',
|
||||
# 'A girl rides a giant cat, walking in the noisy modern city. High definition, realistic, non-cartoonish. Excellent photography work, 8k high definition.', # noqa E501
|
||||
# 11980469406460273604,
|
||||
# ],
|
||||
# [
|
||||
# 'example_inputs/man2.jpeg',
|
||||
# 'example_inputs/woman4.jpeg',
|
||||
# 'ip',
|
||||
# 'ip',
|
||||
# 'a man is dancing with a woman in the room',
|
||||
# 8303780338601106219,
|
||||
# ],
|
||||
# ]
|
||||
# gr.Examples(
|
||||
# examples=example_inps,
|
||||
# inputs=[ref_image1, ref_image2, ref_task1, ref_task2, prompt, seed],
|
||||
# label='Multi IP',
|
||||
# )
|
||||
|
||||
# generate_btn.click(
|
||||
# fn=generate_image,
|
||||
# inputs=[
|
||||
# ref_image1,
|
||||
# ref_image2,
|
||||
# ref_task1,
|
||||
# ref_task2,
|
||||
# prompt,
|
||||
# width,
|
||||
# height,
|
||||
# ref_res,
|
||||
# num_steps,
|
||||
# guidance,
|
||||
# seed,
|
||||
# true_cfg,
|
||||
# cfg_start_step,
|
||||
# cfg_end_step,
|
||||
# neg_prompt,
|
||||
# neg_guidance,
|
||||
# first_step_guidance,
|
||||
# ],
|
||||
# outputs=[output_image, debug_image, seed_output],
|
||||
# )
|
||||
|
||||
# return demo
|
||||
|
||||
|
||||
# if __name__ == '__main__':
|
||||
# demo = create_demo()
|
||||
# demo.queue().launch(server_name='0.0.0.0', server_port=args.port)
|
||||
@@ -0,0 +1,489 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import diffusers
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers import FluxPipeline
|
||||
from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
|
||||
from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
|
||||
from einops import repeat
|
||||
#from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from .transformer import flux_transformer_forward
|
||||
from .utils import convert_flux_lora_to_diffusers
|
||||
|
||||
diffusers.models.transformers.transformer_flux.FluxTransformer2DModel.forward = flux_transformer_forward
|
||||
|
||||
|
||||
def get_task_embedding_idx(task):
|
||||
return 0
|
||||
|
||||
|
||||
class DreamOPipeline(FluxPipeline):
|
||||
def __init__(self, scheduler, vae, text_encoder, tokenizer, text_encoder_2, tokenizer_2, transformer,feature_extractor=None, image_encoder=None):
|
||||
super().__init__(scheduler, vae, text_encoder, tokenizer, text_encoder_2, tokenizer_2, transformer)
|
||||
self.feature_extractor = feature_extractor
|
||||
self.image_encoder = image_encoder
|
||||
self.t5_embedding = nn.Embedding(10, 4096)
|
||||
self.task_embedding = nn.Embedding(2, 3072)
|
||||
self.idx_embedding = nn.Embedding(10, 3072)
|
||||
|
||||
def load_dreamo_model(self, device, dreamo_lora_path,cfg_distill_path,Turbo_path,use_turbo=True):
|
||||
# hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo.safetensors', local_dir='models')
|
||||
# hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo_cfg_distill.safetensors', local_dir='models')
|
||||
dreamo_lora = load_file(dreamo_lora_path)
|
||||
cfg_distill_lora = load_file(cfg_distill_path)
|
||||
self.t5_embedding.weight.data = dreamo_lora.pop('dreamo_t5_embedding.weight')[-10:]
|
||||
self.task_embedding.weight.data = dreamo_lora.pop('dreamo_task_embedding.weight')
|
||||
self.idx_embedding.weight.data = dreamo_lora.pop('dreamo_idx_embedding.weight')
|
||||
self._prepare_t5()
|
||||
|
||||
dreamo_diffuser_lora = convert_flux_lora_to_diffusers(dreamo_lora)
|
||||
cfg_diffuser_lora = convert_flux_lora_to_diffusers(cfg_distill_lora)
|
||||
adapter_names = ['dreamo']
|
||||
adapter_weights = [1]
|
||||
self.load_lora_weights(dreamo_diffuser_lora, adapter_name='dreamo')
|
||||
if cfg_diffuser_lora is not None:
|
||||
self.load_lora_weights(cfg_diffuser_lora, adapter_name='cfg')
|
||||
adapter_names.append('cfg')
|
||||
adapter_weights.append(1)
|
||||
if use_turbo:
|
||||
self.load_lora_weights(
|
||||
Turbo_path,
|
||||
adapter_name='turbo',
|
||||
)
|
||||
adapter_names.append('turbo')
|
||||
adapter_weights.append(1)
|
||||
|
||||
self.fuse_lora(adapter_names=adapter_names, adapter_weights=adapter_weights, lora_scale=1)
|
||||
self.unload_lora_weights()
|
||||
|
||||
self.t5_embedding = self.t5_embedding.to(device)
|
||||
self.task_embedding = self.task_embedding.to(device)
|
||||
self.idx_embedding = self.idx_embedding.to(device)
|
||||
|
||||
def _prepare_t5(self):
|
||||
self.text_encoder_2.resize_token_embeddings(len(self.tokenizer_2))
|
||||
num_new_token = 10
|
||||
new_token_list = [f"[ref#{i}]" for i in range(1, 10)] + ["[res]"]
|
||||
self.tokenizer_2.add_tokens(new_token_list, special_tokens=False)
|
||||
self.text_encoder_2.resize_token_embeddings(len(self.tokenizer_2))
|
||||
input_embedding = self.text_encoder_2.get_input_embeddings().weight.data
|
||||
input_embedding[-num_new_token:] = self.t5_embedding.weight.data
|
||||
|
||||
@staticmethod
|
||||
def _prepare_latent_image_ids(batch_size, height, width, device, dtype, start_height=0, start_width=0):
|
||||
latent_image_ids = torch.zeros(height // 2, width // 2, 3)
|
||||
latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] + start_height
|
||||
latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] + start_width
|
||||
|
||||
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
|
||||
|
||||
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
|
||||
latent_image_ids = latent_image_ids.reshape(
|
||||
batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
|
||||
)
|
||||
|
||||
return latent_image_ids.to(device=device, dtype=dtype)
|
||||
|
||||
# @staticmethod
|
||||
# def _prepare_style_latent_image_ids(batch_size, height, width, device, dtype, start_height=0, start_width=0):
|
||||
# latent_image_ids = torch.zeros(height // 2, width // 2, 3)
|
||||
# latent_image_ids[..., 1] = latent_image_ids[..., 1] + start_height
|
||||
# latent_image_ids[..., 2] = latent_image_ids[..., 2] + start_width
|
||||
|
||||
# latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
|
||||
|
||||
# latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
|
||||
# latent_image_ids = latent_image_ids.reshape(
|
||||
# batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
|
||||
# )
|
||||
|
||||
# return latent_image_ids.to(device=device, dtype=dtype)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _prepare_latent_image_ids(batch_size, height, width, device, dtype, start_height=0, start_width=0):
|
||||
latent_image_ids = torch.zeros(height, width, 3)
|
||||
latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None] + start_height
|
||||
latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :] + start_width
|
||||
|
||||
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
|
||||
|
||||
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
|
||||
latent_image_ids = latent_image_ids.reshape(
|
||||
batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
|
||||
)
|
||||
|
||||
return latent_image_ids.to(device=device, dtype=dtype)
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
negative_prompt: Union[str, List[str]] = None,
|
||||
negative_prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
true_cfg_scale: float = 1.0,
|
||||
true_cfg_start_step: int = 1,
|
||||
true_cfg_end_step: int = 1,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 28,
|
||||
sigmas: Optional[List[float]] = None,
|
||||
guidance_scale: float = 3.5,
|
||||
neg_guidance_scale: float = 3.5,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "latent",
|
||||
return_dict: bool = True,
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 512,
|
||||
ref_conds=None,
|
||||
first_step_guidance_scale=3.5,
|
||||
):
|
||||
r"""
|
||||
Function invoked when calling the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
||||
instead.
|
||||
prompt_2 (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
|
||||
will be used instead.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
|
||||
not greater than `1`).
|
||||
negative_prompt_2 (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
|
||||
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
|
||||
true_cfg_scale (`float`, *optional*, defaults to 1.0):
|
||||
When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
|
||||
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The height in pixels of the generated image. This is set to 1024 by default for the best results.
|
||||
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The width in pixels of the generated image. This is set to 1024 by default for the best results.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
sigmas (`List[float]`, *optional*):
|
||||
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
|
||||
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
|
||||
will be used.
|
||||
guidance_scale (`float`, *optional*, defaults to 3.5):
|
||||
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
||||
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
||||
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
||||
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
||||
usually at the expense of lower image quality.
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
||||
to make generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor will ge generated by sampling using the supplied random `generator`.
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
|
||||
If not provided, pooled text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
|
||||
input argument.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generate image. Choose between
|
||||
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
|
||||
joint_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
callback_on_step_end (`Callable`, *optional*):
|
||||
A function that calls at the end of each denoising steps during the inference. The function is called
|
||||
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
||||
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
||||
`callback_on_step_end_tensor_inputs`.
|
||||
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
||||
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
||||
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
||||
`._callback_tensor_inputs` attribute of your pipeline class.
|
||||
max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
|
||||
is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
|
||||
images.
|
||||
"""
|
||||
|
||||
height = height or self.default_sample_size * self.vae_scale_factor
|
||||
width = width or self.default_sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length=max_sequence_length,
|
||||
)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
self._joint_attention_kwargs = joint_attention_kwargs
|
||||
self._current_timestep = None
|
||||
self._interrupt = False
|
||||
|
||||
# 2. Define call parameters
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
lora_scale = (
|
||||
self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
|
||||
)
|
||||
has_neg_prompt = negative_prompt is not None or (
|
||||
negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
|
||||
)
|
||||
do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
|
||||
(
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
text_ids,
|
||||
) = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
if do_true_cfg:
|
||||
(
|
||||
negative_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
_,
|
||||
) = self.encode_prompt(
|
||||
prompt=negative_prompt,
|
||||
prompt_2=negative_prompt_2,
|
||||
prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
# 4. Prepare latent variables
|
||||
num_channels_latents = self.transformer.config.in_channels // 4
|
||||
latents, latent_image_ids = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 4.1 concat ref tokens to latent
|
||||
origin_img_len = latents.shape[1]
|
||||
embeddings = repeat(self.task_embedding.weight[1], "c -> n l c", n=batch_size, l=origin_img_len)
|
||||
ref_latents = []
|
||||
ref_latent_image_idss = []
|
||||
start_height = height // 16
|
||||
start_width = width // 16
|
||||
for ref_cond in ref_conds:
|
||||
img = ref_cond['img'] # [b, 3, h, w], range [-1, 1]
|
||||
task = ref_cond['task']
|
||||
idx = ref_cond['idx']
|
||||
|
||||
# encode ref with VAE
|
||||
img = img.to(latents)
|
||||
ref_latent = self.vae.encode(img).latent_dist.sample()
|
||||
ref_latent = (ref_latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
|
||||
cur_height = ref_latent.shape[2]
|
||||
cur_width = ref_latent.shape[3]
|
||||
ref_latent = self._pack_latents(ref_latent, batch_size, num_channels_latents, cur_height, cur_width)
|
||||
# ref_latent_image_ids = self._prepare_latent_image_ids(
|
||||
# batch_size, cur_height, cur_width, device, prompt_embeds.dtype, start_height, start_width
|
||||
# )
|
||||
ref_latent_image_ids = self._prepare_latent_image_ids(
|
||||
batch_size, cur_height // 2, cur_width // 2, device, prompt_embeds.dtype, start_height, start_width
|
||||
) #diff 0.33
|
||||
start_height += cur_height // 2
|
||||
start_width += cur_width // 2
|
||||
|
||||
# prepare task_idx_embedding
|
||||
task_idx = get_task_embedding_idx(task) # 0
|
||||
cur_task_embedding = repeat(
|
||||
self.task_embedding.weight[task_idx], "c -> n l c", n=batch_size, l=ref_latent.shape[1]
|
||||
)
|
||||
cur_idx_embedding = repeat(
|
||||
self.idx_embedding.weight[idx], "c -> n l c", n=batch_size, l=ref_latent.shape[1]
|
||||
)
|
||||
cur_embedding = cur_task_embedding + cur_idx_embedding
|
||||
|
||||
# concat ref to latent
|
||||
embeddings = torch.cat([embeddings, cur_embedding], dim=1)
|
||||
ref_latents.append(ref_latent)
|
||||
ref_latent_image_idss.append(ref_latent_image_ids)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
|
||||
image_seq_len = latents.shape[1]
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
self.scheduler.config.get("base_image_seq_len", 256),
|
||||
self.scheduler.config.get("max_image_seq_len", 4096),
|
||||
self.scheduler.config.get("base_shift", 0.5),
|
||||
self.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
mu=mu,
|
||||
)
|
||||
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# handle guidance
|
||||
if self.transformer.config.guidance_embeds:
|
||||
guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
|
||||
guidance = guidance.expand(latents.shape[0])
|
||||
else:
|
||||
guidance = None
|
||||
neg_guidance = torch.full([1], neg_guidance_scale, device=device, dtype=torch.float32)
|
||||
neg_guidance = neg_guidance.expand(latents.shape[0])
|
||||
first_step_guidance = torch.full([1], first_step_guidance_scale, device=device, dtype=torch.float32)
|
||||
|
||||
if self.joint_attention_kwargs is None:
|
||||
self._joint_attention_kwargs = {}
|
||||
|
||||
# 6. Denoising loop
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
if self.interrupt:
|
||||
continue
|
||||
|
||||
self._current_timestep = t
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
||||
#print(latents.shape,pooled_prompt_embeds.shape,prompt_embeds.shape,latent_image_ids.shape,embeddings.shape,text_ids.shape)
|
||||
#[1, 1024, 64],[1, 768],[1, 512, 4096],[1, 256, 3],[1, 2048, 3072],[512,3]
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=torch.cat((latents, *ref_latents), dim=1),
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance if i > 0 else first_step_guidance,
|
||||
pooled_projections=pooled_prompt_embeds,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=torch.cat((latent_image_ids, *ref_latent_image_idss), dim=1),
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
embeddings=embeddings,
|
||||
)[0][:, :origin_img_len]
|
||||
|
||||
if do_true_cfg and i >= true_cfg_start_step and i < true_cfg_end_step:
|
||||
neg_noise_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=neg_guidance,
|
||||
pooled_projections=negative_pooled_prompt_embeds,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=latent_image_ids,
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents_dtype = latents.dtype
|
||||
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
||||
|
||||
if latents.dtype != latents_dtype and torch.backends.mps.is_available():
|
||||
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
|
||||
latents = latents.to(latents_dtype)
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
|
||||
self._current_timestep = None
|
||||
|
||||
if output_type == "latent":
|
||||
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
|
||||
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
|
||||
#print(latents.shape)
|
||||
image = latents
|
||||
else:
|
||||
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
|
||||
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return (image,)
|
||||
|
||||
return FluxPipelineOutput(images=image)
|
||||
@@ -0,0 +1,466 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Any, Callable, Dict, List, Optional, Union
|
||||
|
||||
import diffusers
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers import FluxPipeline
|
||||
from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
|
||||
from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
|
||||
from einops import repeat
|
||||
from huggingface_hub import hf_hub_download
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from .transformer import flux_transformer_forward
|
||||
from .utils import convert_flux_lora_to_diffusers
|
||||
|
||||
diffusers.models.transformers.transformer_flux.FluxTransformer2DModel.forward = flux_transformer_forward
|
||||
|
||||
|
||||
def get_task_embedding_idx(task):
|
||||
return 0
|
||||
|
||||
|
||||
class DreamOPipeline(FluxPipeline):
|
||||
def __init__(self, scheduler, vae, text_encoder, tokenizer, text_encoder_2, tokenizer_2, transformer):
|
||||
super().__init__(scheduler, vae, text_encoder, tokenizer, text_encoder_2, tokenizer_2, transformer)
|
||||
self.t5_embedding = nn.Embedding(10, 4096)
|
||||
self.task_embedding = nn.Embedding(2, 3072)
|
||||
self.idx_embedding = nn.Embedding(10, 3072)
|
||||
|
||||
def load_dreamo_model(self, device, use_turbo=True):
|
||||
hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo.safetensors', local_dir='models')
|
||||
hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo_cfg_distill.safetensors', local_dir='models')
|
||||
dreamo_lora = load_file('models/dreamo.safetensors')
|
||||
cfg_distill_lora = load_file('models/dreamo_cfg_distill.safetensors')
|
||||
self.t5_embedding.weight.data = dreamo_lora.pop('dreamo_t5_embedding.weight')[-10:]
|
||||
self.task_embedding.weight.data = dreamo_lora.pop('dreamo_task_embedding.weight')
|
||||
self.idx_embedding.weight.data = dreamo_lora.pop('dreamo_idx_embedding.weight')
|
||||
self._prepare_t5()
|
||||
|
||||
dreamo_diffuser_lora = convert_flux_lora_to_diffusers(dreamo_lora)
|
||||
cfg_diffuser_lora = convert_flux_lora_to_diffusers(cfg_distill_lora)
|
||||
adapter_names = ['dreamo']
|
||||
adapter_weights = [1]
|
||||
self.load_lora_weights(dreamo_diffuser_lora, adapter_name='dreamo')
|
||||
if cfg_diffuser_lora is not None:
|
||||
self.load_lora_weights(cfg_diffuser_lora, adapter_name='cfg')
|
||||
adapter_names.append('cfg')
|
||||
adapter_weights.append(1)
|
||||
if use_turbo:
|
||||
self.load_lora_weights(
|
||||
hf_hub_download(
|
||||
"alimama-creative/FLUX.1-Turbo-Alpha", "diffusion_pytorch_model.safetensors", local_dir='models'
|
||||
),
|
||||
adapter_name='turbo',
|
||||
)
|
||||
adapter_names.append('turbo')
|
||||
adapter_weights.append(1)
|
||||
|
||||
self.fuse_lora(adapter_names=adapter_names, adapter_weights=adapter_weights, lora_scale=1)
|
||||
|
||||
self.t5_embedding = self.t5_embedding.to(device)
|
||||
self.task_embedding = self.task_embedding.to(device)
|
||||
self.idx_embedding = self.idx_embedding.to(device)
|
||||
|
||||
def _prepare_t5(self):
|
||||
self.text_encoder_2.resize_token_embeddings(len(self.tokenizer_2))
|
||||
num_new_token = 10
|
||||
new_token_list = [f"[ref#{i}]" for i in range(1, 10)] + ["[res]"]
|
||||
self.tokenizer_2.add_tokens(new_token_list, special_tokens=False)
|
||||
self.text_encoder_2.resize_token_embeddings(len(self.tokenizer_2))
|
||||
input_embedding = self.text_encoder_2.get_input_embeddings().weight.data
|
||||
input_embedding[-num_new_token:] = self.t5_embedding.weight.data
|
||||
|
||||
@staticmethod
|
||||
def _prepare_latent_image_ids(batch_size, height, width, device, dtype, start_height=0, start_width=0):
|
||||
latent_image_ids = torch.zeros(height // 2, width // 2, 3)
|
||||
latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] + start_height
|
||||
latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] + start_width
|
||||
|
||||
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
|
||||
|
||||
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
|
||||
latent_image_ids = latent_image_ids.reshape(
|
||||
batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
|
||||
)
|
||||
|
||||
return latent_image_ids.to(device=device, dtype=dtype)
|
||||
|
||||
@staticmethod
|
||||
def _prepare_style_latent_image_ids(batch_size, height, width, device, dtype, start_height=0, start_width=0):
|
||||
latent_image_ids = torch.zeros(height // 2, width // 2, 3)
|
||||
latent_image_ids[..., 1] = latent_image_ids[..., 1] + start_height
|
||||
latent_image_ids[..., 2] = latent_image_ids[..., 2] + start_width
|
||||
|
||||
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
|
||||
|
||||
latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1)
|
||||
latent_image_ids = latent_image_ids.reshape(
|
||||
batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
|
||||
)
|
||||
|
||||
return latent_image_ids.to(device=device, dtype=dtype)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
prompt: Union[str, List[str]] = None,
|
||||
prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
negative_prompt: Union[str, List[str]] = None,
|
||||
negative_prompt_2: Optional[Union[str, List[str]]] = None,
|
||||
true_cfg_scale: float = 1.0,
|
||||
true_cfg_start_step: int = 1,
|
||||
true_cfg_end_step: int = 1,
|
||||
height: Optional[int] = None,
|
||||
width: Optional[int] = None,
|
||||
num_inference_steps: int = 28,
|
||||
sigmas: Optional[List[float]] = None,
|
||||
guidance_scale: float = 3.5,
|
||||
neg_guidance_scale: float = 3.5,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
latents: Optional[torch.FloatTensor] = None,
|
||||
prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
return_dict: bool = True,
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
||||
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
||||
max_sequence_length: int = 512,
|
||||
ref_conds=None,
|
||||
first_step_guidance_scale=3.5,
|
||||
):
|
||||
r"""
|
||||
Function invoked when calling the pipeline for generation.
|
||||
|
||||
Args:
|
||||
prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
||||
instead.
|
||||
prompt_2 (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
|
||||
will be used instead.
|
||||
negative_prompt (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
||||
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
|
||||
not greater than `1`).
|
||||
negative_prompt_2 (`str` or `List[str]`, *optional*):
|
||||
The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
|
||||
`text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
|
||||
true_cfg_scale (`float`, *optional*, defaults to 1.0):
|
||||
When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
|
||||
height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The height in pixels of the generated image. This is set to 1024 by default for the best results.
|
||||
width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
|
||||
The width in pixels of the generated image. This is set to 1024 by default for the best results.
|
||||
num_inference_steps (`int`, *optional*, defaults to 50):
|
||||
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
||||
expense of slower inference.
|
||||
sigmas (`List[float]`, *optional*):
|
||||
Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
|
||||
their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
|
||||
will be used.
|
||||
guidance_scale (`float`, *optional*, defaults to 3.5):
|
||||
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
||||
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
||||
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
||||
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
||||
usually at the expense of lower image quality.
|
||||
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
||||
The number of images to generate per prompt.
|
||||
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
||||
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
||||
to make generation deterministic.
|
||||
latents (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
||||
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
||||
tensor will ge generated by sampling using the supplied random `generator`.
|
||||
prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
||||
provided, text embeddings will be generated from `prompt` input argument.
|
||||
pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
|
||||
If not provided, pooled text embeddings will be generated from `prompt` input argument.
|
||||
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
|
||||
argument.
|
||||
negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
|
||||
Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
|
||||
weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
|
||||
input argument.
|
||||
output_type (`str`, *optional*, defaults to `"pil"`):
|
||||
The output format of the generate image. Choose between
|
||||
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
|
||||
joint_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
callback_on_step_end (`Callable`, *optional*):
|
||||
A function that calls at the end of each denoising steps during the inference. The function is called
|
||||
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
||||
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
||||
`callback_on_step_end_tensor_inputs`.
|
||||
callback_on_step_end_tensor_inputs (`List`, *optional*):
|
||||
The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
|
||||
will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
|
||||
`._callback_tensor_inputs` attribute of your pipeline class.
|
||||
max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
|
||||
|
||||
Examples:
|
||||
|
||||
Returns:
|
||||
[`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
|
||||
is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
|
||||
images.
|
||||
"""
|
||||
|
||||
height = height or self.default_sample_size * self.vae_scale_factor
|
||||
width = width or self.default_sample_size * self.vae_scale_factor
|
||||
|
||||
# 1. Check inputs. Raise error if not correct
|
||||
self.check_inputs(
|
||||
prompt,
|
||||
prompt_2,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
||||
max_sequence_length=max_sequence_length,
|
||||
)
|
||||
|
||||
self._guidance_scale = guidance_scale
|
||||
self._joint_attention_kwargs = joint_attention_kwargs
|
||||
self._current_timestep = None
|
||||
self._interrupt = False
|
||||
|
||||
# 2. Define call parameters
|
||||
if prompt is not None and isinstance(prompt, str):
|
||||
batch_size = 1
|
||||
elif prompt is not None and isinstance(prompt, list):
|
||||
batch_size = len(prompt)
|
||||
else:
|
||||
batch_size = prompt_embeds.shape[0]
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
lora_scale = (
|
||||
self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
|
||||
)
|
||||
has_neg_prompt = negative_prompt is not None or (
|
||||
negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
|
||||
)
|
||||
do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
|
||||
(
|
||||
prompt_embeds,
|
||||
pooled_prompt_embeds,
|
||||
text_ids,
|
||||
) = self.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_2=prompt_2,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
if do_true_cfg:
|
||||
(
|
||||
negative_prompt_embeds,
|
||||
negative_pooled_prompt_embeds,
|
||||
_,
|
||||
) = self.encode_prompt(
|
||||
prompt=negative_prompt,
|
||||
prompt_2=negative_prompt_2,
|
||||
prompt_embeds=negative_prompt_embeds,
|
||||
pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
||||
device=device,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
max_sequence_length=max_sequence_length,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
# 4. Prepare latent variables
|
||||
num_channels_latents = self.transformer.config.in_channels // 4
|
||||
latents, latent_image_ids = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
height,
|
||||
width,
|
||||
prompt_embeds.dtype,
|
||||
device,
|
||||
generator,
|
||||
latents,
|
||||
)
|
||||
|
||||
# 4.1 concat ref tokens to latent
|
||||
origin_img_len = latents.shape[1]
|
||||
embeddings = repeat(self.task_embedding.weight[1], "c -> n l c", n=batch_size, l=origin_img_len)
|
||||
ref_latents = []
|
||||
ref_latent_image_idss = []
|
||||
start_height = height // 16
|
||||
start_width = width // 16
|
||||
for ref_cond in ref_conds:
|
||||
img = ref_cond['img'] # [b, 3, h, w], range [-1, 1]
|
||||
task = ref_cond['task']
|
||||
idx = ref_cond['idx']
|
||||
|
||||
# encode ref with VAE
|
||||
img = img.to(latents)
|
||||
ref_latent = self.vae.encode(img).latent_dist.sample()
|
||||
ref_latent = (ref_latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
|
||||
cur_height = ref_latent.shape[2]
|
||||
cur_width = ref_latent.shape[3]
|
||||
ref_latent = self._pack_latents(ref_latent, batch_size, num_channels_latents, cur_height, cur_width)
|
||||
ref_latent_image_ids = self._prepare_latent_image_ids(
|
||||
batch_size, cur_height, cur_width, device, prompt_embeds.dtype, start_height, start_width
|
||||
)
|
||||
start_height += cur_height // 2
|
||||
start_width += cur_width // 2
|
||||
|
||||
# prepare task_idx_embedding
|
||||
task_idx = get_task_embedding_idx(task)
|
||||
cur_task_embedding = repeat(
|
||||
self.task_embedding.weight[task_idx], "c -> n l c", n=batch_size, l=ref_latent.shape[1]
|
||||
)
|
||||
cur_idx_embedding = repeat(
|
||||
self.idx_embedding.weight[idx], "c -> n l c", n=batch_size, l=ref_latent.shape[1]
|
||||
)
|
||||
cur_embedding = cur_task_embedding + cur_idx_embedding
|
||||
|
||||
# concat ref to latent
|
||||
embeddings = torch.cat([embeddings, cur_embedding], dim=1)
|
||||
ref_latents.append(ref_latent)
|
||||
ref_latent_image_idss.append(ref_latent_image_ids)
|
||||
|
||||
# 5. Prepare timesteps
|
||||
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
|
||||
image_seq_len = latents.shape[1]
|
||||
mu = calculate_shift(
|
||||
image_seq_len,
|
||||
self.scheduler.config.get("base_image_seq_len", 256),
|
||||
self.scheduler.config.get("max_image_seq_len", 4096),
|
||||
self.scheduler.config.get("base_shift", 0.5),
|
||||
self.scheduler.config.get("max_shift", 1.15),
|
||||
)
|
||||
timesteps, num_inference_steps = retrieve_timesteps(
|
||||
self.scheduler,
|
||||
num_inference_steps,
|
||||
device,
|
||||
sigmas=sigmas,
|
||||
mu=mu,
|
||||
)
|
||||
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
||||
self._num_timesteps = len(timesteps)
|
||||
|
||||
# handle guidance
|
||||
if self.transformer.config.guidance_embeds:
|
||||
guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
|
||||
guidance = guidance.expand(latents.shape[0])
|
||||
else:
|
||||
guidance = None
|
||||
neg_guidance = torch.full([1], neg_guidance_scale, device=device, dtype=torch.float32)
|
||||
neg_guidance = neg_guidance.expand(latents.shape[0])
|
||||
first_step_guidance = torch.full([1], first_step_guidance_scale, device=device, dtype=torch.float32)
|
||||
|
||||
if self.joint_attention_kwargs is None:
|
||||
self._joint_attention_kwargs = {}
|
||||
|
||||
# 6. Denoising loop
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
if self.interrupt:
|
||||
continue
|
||||
|
||||
self._current_timestep = t
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timestep = t.expand(latents.shape[0]).to(latents.dtype)
|
||||
|
||||
noise_pred = self.transformer(
|
||||
hidden_states=torch.cat((latents, *ref_latents), dim=1),
|
||||
timestep=timestep / 1000,
|
||||
guidance=guidance if i > 0 else first_step_guidance,
|
||||
pooled_projections=pooled_prompt_embeds,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=torch.cat((latent_image_ids, *ref_latent_image_idss), dim=1),
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
embeddings=embeddings,
|
||||
)[0][:, :origin_img_len]
|
||||
|
||||
if do_true_cfg and i >= true_cfg_start_step and i < true_cfg_end_step:
|
||||
neg_noise_pred = self.transformer(
|
||||
hidden_states=latents,
|
||||
timestep=timestep / 1000,
|
||||
guidance=neg_guidance,
|
||||
pooled_projections=negative_pooled_prompt_embeds,
|
||||
encoder_hidden_states=negative_prompt_embeds,
|
||||
txt_ids=text_ids,
|
||||
img_ids=latent_image_ids,
|
||||
joint_attention_kwargs=self.joint_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents_dtype = latents.dtype
|
||||
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
|
||||
|
||||
if latents.dtype != latents_dtype and torch.backends.mps.is_available():
|
||||
# some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
|
||||
latents = latents.to(latents_dtype)
|
||||
|
||||
if callback_on_step_end is not None:
|
||||
callback_kwargs = {}
|
||||
for k in callback_on_step_end_tensor_inputs:
|
||||
callback_kwargs[k] = locals()[k]
|
||||
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
||||
|
||||
latents = callback_outputs.pop("latents", latents)
|
||||
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
||||
progress_bar.update()
|
||||
|
||||
self._current_timestep = None
|
||||
|
||||
if output_type == "latent":
|
||||
image = latents
|
||||
else:
|
||||
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
|
||||
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
|
||||
image = self.vae.decode(latents, return_dict=False)[0]
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
|
||||
# Offload all models
|
||||
self.maybe_free_model_hooks()
|
||||
|
||||
if not return_dict:
|
||||
return (image,)
|
||||
|
||||
return FluxPipelineOutput(images=image)
|
||||
@@ -0,0 +1,187 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.utils import (
|
||||
USE_PEFT_BACKEND,
|
||||
logging,
|
||||
scale_lora_layers,
|
||||
unscale_lora_layers,
|
||||
)
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def flux_transformer_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor = None,
|
||||
pooled_projections: torch.Tensor = None,
|
||||
timestep: torch.LongTensor = None,
|
||||
img_ids: torch.Tensor = None,
|
||||
txt_ids: torch.Tensor = None,
|
||||
guidance: torch.Tensor = None,
|
||||
joint_attention_kwargs: Optional[Dict[str, Any]] = None,
|
||||
controlnet_block_samples=None,
|
||||
controlnet_single_block_samples=None,
|
||||
return_dict: bool = True,
|
||||
controlnet_blocks_repeat: bool = False,
|
||||
embeddings: torch.Tensor = None,
|
||||
) -> Union[torch.Tensor, Transformer2DModelOutput]:
|
||||
"""
|
||||
The [`FluxTransformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
|
||||
Input `hidden_states`.
|
||||
encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
|
||||
Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
|
||||
pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
|
||||
from the embeddings of input conditions.
|
||||
timestep ( `torch.LongTensor`):
|
||||
Used to indicate denoising step.
|
||||
block_controlnet_hidden_states: (`list` of `torch.Tensor`):
|
||||
A list of tensors that if specified are added to the residuals of transformer blocks.
|
||||
joint_attention_kwargs (`dict`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
if joint_attention_kwargs is not None:
|
||||
joint_attention_kwargs = joint_attention_kwargs.copy()
|
||||
lora_scale = joint_attention_kwargs.pop("scale", 1.0)
|
||||
else:
|
||||
lora_scale = 1.0
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# weight the lora layers by setting `lora_scale` for each PEFT layer
|
||||
scale_lora_layers(self, lora_scale)
|
||||
else:
|
||||
if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
|
||||
logger.warning(
|
||||
"Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
|
||||
)
|
||||
|
||||
hidden_states = self.x_embedder(hidden_states)
|
||||
# add task and idx embedding
|
||||
if embeddings is not None:
|
||||
hidden_states = hidden_states + embeddings
|
||||
|
||||
timestep = timestep.to(hidden_states.dtype) * 1000
|
||||
guidance = guidance.to(hidden_states.dtype) * 1000 if guidance is not None else None
|
||||
|
||||
temb = (
|
||||
self.time_text_embed(timestep, pooled_projections)
|
||||
if guidance is None
|
||||
else self.time_text_embed(timestep, guidance, pooled_projections)
|
||||
)
|
||||
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
|
||||
|
||||
if txt_ids.ndim == 3:
|
||||
# logger.warning(
|
||||
# "Passing `txt_ids` 3d torch.Tensor is deprecated."
|
||||
# "Please remove the batch dimension and pass it as a 2d torch Tensor"
|
||||
# )
|
||||
txt_ids = txt_ids[0]
|
||||
if img_ids.ndim == 3:
|
||||
# logger.warning(
|
||||
# "Passing `img_ids` 3d torch.Tensor is deprecated."
|
||||
# "Please remove the batch dimension and pass it as a 2d torch Tensor"
|
||||
# )
|
||||
img_ids = img_ids[0]
|
||||
|
||||
ids = torch.cat((txt_ids, img_ids), dim=0)
|
||||
image_rotary_emb = self.pos_embed(ids)
|
||||
|
||||
for index_block, block in enumerate(self.transformer_blocks):
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
|
||||
block,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
)
|
||||
|
||||
else:
|
||||
encoder_hidden_states, hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
joint_attention_kwargs=joint_attention_kwargs,
|
||||
)
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_block_samples is not None:
|
||||
interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
# For Xlabs ControlNet.
|
||||
if controlnet_blocks_repeat:
|
||||
hidden_states = hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
|
||||
else:
|
||||
hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
|
||||
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
||||
|
||||
for index_block, block in enumerate(self.single_transformer_blocks):
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
hidden_states = self._gradient_checkpointing_func(
|
||||
block,
|
||||
hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
)
|
||||
|
||||
else:
|
||||
hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
joint_attention_kwargs=joint_attention_kwargs,
|
||||
)
|
||||
|
||||
# controlnet residual
|
||||
if controlnet_single_block_samples is not None:
|
||||
interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
|
||||
interval_control = int(np.ceil(interval_control))
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
|
||||
hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
+ controlnet_single_block_samples[index_block // interval_control]
|
||||
)
|
||||
|
||||
hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
|
||||
|
||||
hidden_states = self.norm_out(hidden_states, temb)
|
||||
output = self.proj_out(hidden_states)
|
||||
|
||||
if USE_PEFT_BACKEND:
|
||||
# remove `lora_scale` from each PEFT layer
|
||||
unscale_lora_layers(self, lora_scale)
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
@@ -0,0 +1,222 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
import re
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.utils import make_grid
|
||||
|
||||
|
||||
# from basicsr
|
||||
def img2tensor(imgs, bgr2rgb=True, float32=True):
|
||||
"""Numpy array to tensor.
|
||||
|
||||
Args:
|
||||
imgs (list[ndarray] | ndarray): Input images.
|
||||
bgr2rgb (bool): Whether to change bgr to rgb.
|
||||
float32 (bool): Whether to change to float32.
|
||||
|
||||
Returns:
|
||||
list[tensor] | tensor: Tensor images. If returned results only have
|
||||
one element, just return tensor.
|
||||
"""
|
||||
|
||||
def _totensor(img, bgr2rgb, float32):
|
||||
if img.shape[2] == 3 and bgr2rgb:
|
||||
if img.dtype == 'float64':
|
||||
img = img.astype('float32')
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = torch.from_numpy(img.transpose(2, 0, 1))
|
||||
if float32:
|
||||
img = img.float()
|
||||
return img
|
||||
|
||||
if isinstance(imgs, list):
|
||||
return [_totensor(img, bgr2rgb, float32) for img in imgs]
|
||||
return _totensor(imgs, bgr2rgb, float32)
|
||||
|
||||
|
||||
def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
|
||||
"""Convert torch Tensors into image numpy arrays.
|
||||
|
||||
After clamping to [min, max], values will be normalized to [0, 1].
|
||||
|
||||
Args:
|
||||
tensor (Tensor or list[Tensor]): Accept shapes:
|
||||
1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
|
||||
2) 3D Tensor of shape (3/1 x H x W);
|
||||
3) 2D Tensor of shape (H x W).
|
||||
Tensor channel should be in RGB order.
|
||||
rgb2bgr (bool): Whether to change rgb to bgr.
|
||||
out_type (numpy type): output types. If ``np.uint8``, transform outputs
|
||||
to uint8 type with range [0, 255]; otherwise, float type with
|
||||
range [0, 1]. Default: ``np.uint8``.
|
||||
min_max (tuple[int]): min and max values for clamp.
|
||||
|
||||
Returns:
|
||||
(Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
|
||||
shape (H x W). The channel order is BGR.
|
||||
"""
|
||||
if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
|
||||
raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
|
||||
|
||||
if torch.is_tensor(tensor):
|
||||
tensor = [tensor]
|
||||
result = []
|
||||
for _tensor in tensor:
|
||||
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
|
||||
_tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
|
||||
|
||||
n_dim = _tensor.dim()
|
||||
if n_dim == 4:
|
||||
img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
|
||||
img_np = img_np.transpose(1, 2, 0)
|
||||
if rgb2bgr:
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
elif n_dim == 3:
|
||||
img_np = _tensor.numpy()
|
||||
img_np = img_np.transpose(1, 2, 0)
|
||||
if img_np.shape[2] == 1: # gray image
|
||||
img_np = np.squeeze(img_np, axis=2)
|
||||
else:
|
||||
if rgb2bgr:
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
elif n_dim == 2:
|
||||
img_np = _tensor.numpy()
|
||||
else:
|
||||
raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
|
||||
if out_type == np.uint8:
|
||||
# Unlike MATLAB, numpy.unit8() WILL NOT round by default.
|
||||
img_np = (img_np * 255.0).round()
|
||||
img_np = img_np.astype(out_type)
|
||||
result.append(img_np)
|
||||
if len(result) == 1:
|
||||
result = result[0]
|
||||
return result
|
||||
|
||||
|
||||
def resize_numpy_image_area(image, area=512 * 512):
|
||||
h, w = image.shape[:2]
|
||||
k = math.sqrt(area / (h * w))
|
||||
h = int(h * k) - (int(h * k) % 16)
|
||||
w = int(w * k) - (int(w * k) % 16)
|
||||
image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
|
||||
return image
|
||||
|
||||
|
||||
# reference: https://github.com/huggingface/diffusers/pull/9295/files
|
||||
def convert_flux_lora_to_diffusers(old_state_dict):
|
||||
new_state_dict = {}
|
||||
orig_keys = list(old_state_dict.keys())
|
||||
|
||||
def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
|
||||
down_weight = sds_sd.pop(sds_key)
|
||||
up_weight = sds_sd.pop(sds_key.replace(".down.weight", ".up.weight"))
|
||||
|
||||
# calculate dims if not provided
|
||||
num_splits = len(ait_keys)
|
||||
if dims is None:
|
||||
dims = [up_weight.shape[0] // num_splits] * num_splits
|
||||
else:
|
||||
assert sum(dims) == up_weight.shape[0]
|
||||
|
||||
# make ai-toolkit weight
|
||||
ait_down_keys = [k + ".lora_A.weight" for k in ait_keys]
|
||||
ait_up_keys = [k + ".lora_B.weight" for k in ait_keys]
|
||||
|
||||
# down_weight is copied to each split
|
||||
ait_sd.update({k: down_weight for k in ait_down_keys})
|
||||
|
||||
# up_weight is split to each split
|
||||
ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))}) # noqa: C416
|
||||
|
||||
for old_key in orig_keys:
|
||||
# Handle double_blocks
|
||||
if 'double_blocks' in old_key:
|
||||
block_num = re.search(r"double_blocks_(\d+)", old_key).group(1)
|
||||
new_key = f"transformer.transformer_blocks.{block_num}"
|
||||
|
||||
if "proj_lora1" in old_key:
|
||||
new_key += ".attn.to_out.0"
|
||||
elif "proj_lora2" in old_key:
|
||||
new_key += ".attn.to_add_out"
|
||||
elif "qkv_lora2" in old_key and "up" not in old_key:
|
||||
handle_qkv(
|
||||
old_state_dict,
|
||||
new_state_dict,
|
||||
old_key,
|
||||
[
|
||||
f"transformer.transformer_blocks.{block_num}.attn.add_q_proj",
|
||||
f"transformer.transformer_blocks.{block_num}.attn.add_k_proj",
|
||||
f"transformer.transformer_blocks.{block_num}.attn.add_v_proj",
|
||||
],
|
||||
)
|
||||
# continue
|
||||
elif "qkv_lora1" in old_key and "up" not in old_key:
|
||||
handle_qkv(
|
||||
old_state_dict,
|
||||
new_state_dict,
|
||||
old_key,
|
||||
[
|
||||
f"transformer.transformer_blocks.{block_num}.attn.to_q",
|
||||
f"transformer.transformer_blocks.{block_num}.attn.to_k",
|
||||
f"transformer.transformer_blocks.{block_num}.attn.to_v",
|
||||
],
|
||||
)
|
||||
# continue
|
||||
|
||||
if "down" in old_key:
|
||||
new_key += ".lora_A.weight"
|
||||
elif "up" in old_key:
|
||||
new_key += ".lora_B.weight"
|
||||
|
||||
# Handle single_blocks
|
||||
elif 'single_blocks' in old_key:
|
||||
block_num = re.search(r"single_blocks_(\d+)", old_key).group(1)
|
||||
new_key = f"transformer.single_transformer_blocks.{block_num}"
|
||||
|
||||
if "proj_lora" in old_key:
|
||||
new_key += ".proj_out"
|
||||
elif "qkv_lora" in old_key and "up" not in old_key:
|
||||
handle_qkv(
|
||||
old_state_dict,
|
||||
new_state_dict,
|
||||
old_key,
|
||||
[
|
||||
f"transformer.single_transformer_blocks.{block_num}.attn.to_q",
|
||||
f"transformer.single_transformer_blocks.{block_num}.attn.to_k",
|
||||
f"transformer.single_transformer_blocks.{block_num}.attn.to_v",
|
||||
],
|
||||
)
|
||||
|
||||
if "down" in old_key:
|
||||
new_key += ".lora_A.weight"
|
||||
elif "up" in old_key:
|
||||
new_key += ".lora_B.weight"
|
||||
|
||||
else:
|
||||
# Handle other potential key patterns here
|
||||
new_key = old_key
|
||||
|
||||
# Since we already handle qkv above.
|
||||
if "qkv" not in old_key and 'embedding' not in old_key:
|
||||
new_state_dict[new_key] = old_state_dict.pop(old_key)
|
||||
|
||||
# if len(old_state_dict) > 0:
|
||||
# raise ValueError(f"`old_state_dict` should be at this point but has: {list(old_state_dict.keys())}.")
|
||||
|
||||
return new_state_dict
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 54 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 658 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 585 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 676 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 179 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 752 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 103 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 55 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 414 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 842 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 186 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 737 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 144 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 92 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 447 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 231 KiB |
@@ -0,0 +1,29 @@
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
exclude = ['tools']
|
||||
# A list of file patterns to omit from linting, in addition to those specified by exclude.
|
||||
extend-exclude = ["__pycache__", "*.pyc", "*.egg-info", ".cache"]
|
||||
|
||||
select = ["E", "F", "W", "C90", "I", "UP", "B", "C4", "RET", "RUF", "SIM"]
|
||||
|
||||
|
||||
ignore = [
|
||||
"UP006", # UP006: Use list instead of typing.List for type annotations
|
||||
"UP007", # UP007: Use X | Y for type annotations
|
||||
"UP009",
|
||||
"UP035",
|
||||
"UP038",
|
||||
"E402",
|
||||
"RET504",
|
||||
"C901",
|
||||
"RUF013",
|
||||
"B006",
|
||||
]
|
||||
|
||||
[tool.isort]
|
||||
profile = "black"
|
||||
|
||||
[tool.black]
|
||||
line-length = 119
|
||||
skip-string-normalization = 1
|
||||
exclude = 'tools'
|
||||
@@ -0,0 +1,15 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cu118
|
||||
torch==2.3.1+cu118
|
||||
torchvision==0.18.1+cu118
|
||||
|
||||
einops
|
||||
timm
|
||||
diffusers==0.31.0
|
||||
transformers==4.45.2
|
||||
sentencepiece
|
||||
gradio==4.44.1
|
||||
spaces
|
||||
huggingface_hub
|
||||
accelerate==0.32.0
|
||||
peft
|
||||
git+https://github.com/ToTheBeginning/facexlib.git
|
||||
File diff suppressed because it is too large
Load Diff
+115
-9
@@ -8,7 +8,8 @@ import os
|
||||
from PIL import ImageFont,Image
|
||||
import torch.nn.functional as F
|
||||
import copy
|
||||
|
||||
from pathlib import PureWindowsPath
|
||||
from tqdm import tqdm
|
||||
from .utils.utils import get_comic
|
||||
from .model_loader_utils import (phi2narry,replicate_data_by_indices,get_float,gc_cleanup,tensor_to_image,photomaker_clip,tensortopil_list_upscale,tensortopil_list,extract_content_from_brackets_,
|
||||
narry_list_pil,pre_text2infer,cf_clip,get_phrases_idx_cf,get_eot_idx_cf,get_ms_phrase_emb,get_extra_function,photomaker_clip_v2,adjust_indices,load_clip_clipvsion,
|
||||
@@ -64,6 +65,12 @@ class EasyFunction_Lite:
|
||||
CATEGORY = "Storydiffusion"
|
||||
|
||||
def easy_function_main(self, extra_repo,checkpoints,clip,clip_vision,lora,function_mode,select_method,temperature):
|
||||
|
||||
if extra_repo:
|
||||
extra_repo=PureWindowsPath(extra_repo).as_posix()
|
||||
if select_method:
|
||||
select_method=PureWindowsPath(select_method).as_posix()
|
||||
|
||||
use_gguf,use_unet=False,False
|
||||
|
||||
use_svdq=True if "svdq" in select_method else False
|
||||
@@ -113,11 +120,12 @@ class EasyFunction_Lite:
|
||||
add_function["clip_path"]=clip_path
|
||||
add_function["clip_vision_path"]=clip_vision_path
|
||||
add_function["clip_path_dual"]=clip_path_dual
|
||||
add_function["clip_vision_path_dual"]=clip_vision_path_dual
|
||||
add_function["clip_vision_path_dual"]=clip_vision_path_dual
|
||||
|
||||
else:
|
||||
model=None
|
||||
|
||||
pipe={"model":model,"extra_repo":extra_repo,"model_path":model_path,"use_svdq":use_svdq,"use_gguf":use_gguf,"clip_vision_path_dual":clip_vision_path_dual,
|
||||
pipe={"model":model,"extra_repo":extra_repo,"model_path":model_path,"use_svdq":use_svdq,"use_gguf":use_gguf,"clip_vision_path_dual":clip_vision_path_dual,"lora_dual_path":lora_dual_path,
|
||||
"lora_ckpt_path":lora_path,"use_unet":use_unet,"extra_easy":function_mode,"select_method":select_method,"clip_vision_path":clip_vision_path}
|
||||
return (pipe,clip,add_function)
|
||||
|
||||
@@ -129,7 +137,7 @@ class StoryDiffusion_Apply:
|
||||
{
|
||||
"model": ("MODEL",),
|
||||
"vae": ("VAE",),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","realcustom","instant_character","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"infer_mode": (["story", "classic","flux_pulid","infiniteyou","uno","realcustom","instant_character","dreamo","story_maker","story_and_maker","consistory","kolor_face","msdiffusion" ],),
|
||||
"photomake_ckpt": (["none"] + [i for i in folder_paths.get_filename_list("photomaker") if "v1" in i or "v2" in i],),
|
||||
"ipadapter_ckpt": (["none"] + folder_paths.get_filename_list("photomaker"),),
|
||||
"quantize_mode": ([ "fp8", "nf4","fp16", ],),
|
||||
@@ -151,6 +159,10 @@ class StoryDiffusion_Apply:
|
||||
# pre data
|
||||
photomake_ckpt_path = None if photomake_ckpt == "none" else folder_paths.get_full_path("photomaker", photomake_ckpt)
|
||||
ipadapter_ckpt_path = None if ipadapter_ckpt == "none" else folder_paths.get_full_path("photomaker", ipadapter_ckpt)
|
||||
|
||||
if extra_function:
|
||||
extra_function=PureWindowsPath(extra_function).as_posix()
|
||||
|
||||
save_quantezed=True if "save" in extra_function and quantize_mode=="fp8" else False
|
||||
clip_vision_path=model.get("clip_vision_path") if isinstance(model,dict) else None
|
||||
clip_vision_path_dual=model.get("clip_vision_path_dual") if isinstance(model,dict) else None
|
||||
@@ -172,7 +184,20 @@ class StoryDiffusion_Apply:
|
||||
from comfy.clip_vision import load as clip_load
|
||||
CLIP_VISION=clip_load(clip_vision_path).model
|
||||
|
||||
|
||||
|
||||
dreamo_lora=model.get("lora_ckpt_path") if isinstance(model,dict) else None
|
||||
cfg_distill=model.get("lora_dual_path") if isinstance(model,dict) else None
|
||||
if infer_mode =="dreamo" and dreamo_lora is not None and cfg_distill is not None:
|
||||
if "distill" in dreamo_lora:
|
||||
cfg_distill_path=dreamo_lora
|
||||
dreamo_lora_path=cfg_distill
|
||||
else:
|
||||
cfg_distill_path=cfg_distill
|
||||
dreamo_lora_path=dreamo_lora
|
||||
else:
|
||||
cfg_distill_path=None
|
||||
dreamo_lora_path=None
|
||||
|
||||
|
||||
if infer_mode in ["story_maker" ,"story_and_maker"] and not CLIP_VISION and ipadapter_ckpt_path is None:
|
||||
raise "story_maker need a clipvison H model,mask.bin"
|
||||
@@ -238,6 +263,11 @@ class StoryDiffusion_Apply:
|
||||
elif infer_mode == "uno":
|
||||
from .model_loader_utils import Loader_UNO
|
||||
model = Loader_UNO(model,offload,quantize_mode,save_quantezed,lora_rank=512)
|
||||
elif infer_mode == "dreamo":
|
||||
from.model_loader_utils import Loader_Dreamo
|
||||
if dreamo_lora_path is None or cfg_distill_path is None or ipadapter_ckpt_path is None:
|
||||
raise "dreamo need a dreamo lora and cfg distill and turbo lora in ipadapter menu"
|
||||
model = Loader_Dreamo(model,vae,quantize_mode,dreamo_lora_path,cfg_distill_path,ipadapter_ckpt_path,device)
|
||||
else: # can not choice a mode
|
||||
print("infer use comfyui classic mode")
|
||||
|
||||
@@ -295,7 +325,8 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
add_function=kwargs.get("add_function")
|
||||
image=kwargs.get("image")
|
||||
control_image=kwargs.get("control_image")
|
||||
|
||||
if extra_function:
|
||||
extra_function=PureWindowsPath(extra_function).as_posix()
|
||||
tag_dict,clip_path,text_model,vision_model,clip_path_dual,siglip_path,dino_path=None,None,None,None,None,None,None
|
||||
|
||||
|
||||
@@ -434,10 +465,20 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
|
||||
# get emb use insightface
|
||||
pass
|
||||
elif infer_mode == "infiniteyou":
|
||||
pass
|
||||
elif infer_mode == "dreamo":
|
||||
from .model_loader_utils import Dreamo_image_encoder
|
||||
from huggingface_hub import hf_hub_download
|
||||
BEN2_path=switch.get("select_method")
|
||||
if not BEN2_path:
|
||||
BEN2_path= hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='ComfyUI/models')
|
||||
ref_list=tensortopil_list_upscale(image,width,height)
|
||||
task_list=["ip", "id", "style"] #TODO
|
||||
image_emb={}
|
||||
for key,role_img in zip(role_list,ref_list):
|
||||
role_emb=Dreamo_image_encoder(BEN2_path,role_img,None,task_list[0],task_list[0],ref_res=512) #TODO
|
||||
image_emb[key]=role_emb
|
||||
|
||||
elif infer_mode == "realcustom":
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
text_model,vision_model=load_clip_clipvsion([clip_path,clip_path_dual],
|
||||
[os.path.join(dir_path, "config/clip_1"),os.path.join(dir_path, "config/clip_2")],
|
||||
@@ -653,6 +694,8 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
else:
|
||||
if infer_mode=="classic": #TODO 逆序的角色会出现iD不匹配，受影响的有story文生图
|
||||
only_role_emb= cf_clip(only_role_list, clip, infer_mode,role_list) #story模式需要拆分prompt，所以这里需要传入role_list
|
||||
elif infer_mode=="dreamo":
|
||||
pass # TODO 暂时不支持dreamo
|
||||
else:
|
||||
only_role_emb= cf_clip(inf_list_split, clip, infer_mode,role_list) #story,story_maker,story_and_maker,msdiffusion,infinite
|
||||
# pre nc txt emb
|
||||
@@ -670,6 +713,8 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
else:
|
||||
if infer_mode!="kolor_face":
|
||||
nc_emb=cf_clip(nc_txt_list, clip, infer_mode,role_list,input_split=False)
|
||||
elif infer_mode=="dreamo":
|
||||
pass # TODO 暂时不支持dreamo
|
||||
else:
|
||||
nc_emb,_= glm_single_encode(chatglm3_model, nc_txt_list,role_list, neg_text, 1,nc_mode=True)
|
||||
else:
|
||||
@@ -743,6 +788,19 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
for dual_t,x_1 in zip(prompts_dual,x_1_refs_dual): # dual_t:The figurine play whith The pig in the garden best 8k,RAW
|
||||
inp = prepare_multi_ip_wrapper(clip,prompt=dual_t, ref_imgs=x_1, pe=uno_pe,device=device,h=height,w=width)
|
||||
daul_emb.append(inp)
|
||||
elif prompts_dual and infer_mode == "dreamo":
|
||||
from .model_loader_utils import Dreamo_image_encoder
|
||||
from huggingface_hub import hf_hub_download
|
||||
BEN2_path=switch.get("select_method")
|
||||
if not BEN2_path:
|
||||
BEN2_path= hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='ComfyUI/models')
|
||||
ref_list=tensortopil_list_upscale(image,width,height)
|
||||
task_list=["ip", "id", "style"] #TODO
|
||||
images_emb=Dreamo_image_encoder(BEN2_path,ref_list[0],ref_list[1],task_list[0],task_list[0],ref_res=512) #TODO
|
||||
prompts_dual=[i.replace(role_list[0] ,role_dict[role_list[0]]) for i in prompts_dual if role_list[0] in i ]
|
||||
prompts_dual = [i.replace(role_list[1], role_dict[role_list[1]]) for i in prompts_dual if role_list[1] in i]
|
||||
prompts_dual=[apply_style_positive(add_style,i+pos_text)[0] for i in prompts_dual] #[' The figurine play whith The pig in the garden,best 8k,RAW']
|
||||
daul_emb=[images_emb,prompts_dual]
|
||||
else:
|
||||
daul_emb=None
|
||||
# neg
|
||||
@@ -752,6 +810,12 @@ class StoryDiffusion_CLIPTextEncode:
|
||||
elif infer_mode=="instant_character":
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb}
|
||||
negative = None
|
||||
elif infer_mode=="dreamo":
|
||||
only_role_emb={}
|
||||
for key ,prompts in zip(role_list,inf_list_split):
|
||||
only_role_emb[key]=prompts
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb}
|
||||
negative = neg_text # TODO
|
||||
elif infer_mode=="realcustom":
|
||||
postive_dict= {"role": only_role_emb, "nc": None, "daul": daul_emb}
|
||||
negative=[letent_real]
|
||||
@@ -1508,10 +1572,52 @@ class StoryDiffusion_KSampler:
|
||||
)[0] # torch.Size([1, 4, 64, 64])
|
||||
#print(samples.shape)
|
||||
samples_list.append(samples)
|
||||
|
||||
out = {}
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
elif infer_mode=="dreamo":
|
||||
samples_list = []
|
||||
|
||||
first_step_guidance=0
|
||||
for key in role_list:
|
||||
for prompt in tqdm(only_role_emb[key],desc=f"Processing {key}"):
|
||||
samples = model(prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
ref_conds=image_embeds[key],
|
||||
generator=torch.Generator(device="cpu").manual_seed(seed),
|
||||
true_cfg_scale=1,
|
||||
true_cfg_start_step=0,
|
||||
true_cfg_end_step=0,
|
||||
negative_prompt=neg_text,
|
||||
neg_guidance_scale=cfg,
|
||||
first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else cfg,
|
||||
).images
|
||||
#print(samples.shape)
|
||||
samples_list.append(samples)
|
||||
if daul_emb:
|
||||
for index, prompt in tqdm(zip(dual_index, daul_emb[1]),desc="Processing dual prompts"): #daul_emb [emb,prompts]
|
||||
samples = model(prompt=prompt,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=cfg,
|
||||
ref_conds=daul_emb[0],
|
||||
generator=torch.Generator(device="cpu").manual_seed(seed),
|
||||
true_cfg_scale=1,
|
||||
true_cfg_start_step=0,
|
||||
true_cfg_end_step=0,
|
||||
negative_prompt=neg_text,
|
||||
neg_guidance_scale=cfg,
|
||||
first_step_guidance_scale=first_step_guidance if first_step_guidance > 0 else cfg,
|
||||
).images
|
||||
samples_list.insert(index, samples)
|
||||
out = {}
|
||||
out["samples"] = torch.cat(samples_list, dim=0)
|
||||
return (out,)
|
||||
else: #none:
|
||||
return
|
||||
|
||||
|
||||
+1206
-584
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 366 KiB |
@@ -2776,3 +2776,144 @@ def instant_character_id_clip(subject_image,siglip_image_encoder,siglip_image_pr
|
||||
subject_image_embeds_dict = encode_image_emb(subject_image, device, dtype)
|
||||
|
||||
return subject_image_embeds_dict
|
||||
|
||||
def Loader_Dreamo(cf_model,VAE,quantize_mode,dreamo_lora_path,cfg_distill_path,Turbo_path,device):
|
||||
from .DreamO.dreamo.dreamo_pipeline import DreamOPipeline
|
||||
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig,FluxTransformer2DModel
|
||||
vae=convert_cfvae2diffuser(VAE,use_flux=True)
|
||||
if cf_model.get("use_svdq") :
|
||||
print("use svdq")
|
||||
from nunchaku import NunchakuFluxTransformer2dModel
|
||||
transformer = NunchakuFluxTransformer2dModel.from_pretrained(cf_model.get("select_method"),offload=True)
|
||||
try:
|
||||
transformer.set_attention_impl("nunchaku-fp16")
|
||||
except:
|
||||
pass
|
||||
dreamo_pipeline = DreamOPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"), vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
|
||||
elif cf_model.get("use_gguf"):
|
||||
print("use gguf quantization")
|
||||
if cf_model.get("model_path") is None:
|
||||
raise "need chocie a gguf model in EasyFunction_Lite!"
|
||||
from diffusers import GGUFQuantizationConfig
|
||||
transformer = FluxTransformer2DModel.from_single_file(
|
||||
cf_model.get("model_path"),
|
||||
config=os.path.join(cur_path, "config/FLUX.1-dev/transformer"),
|
||||
quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
vae=convert_cfvae2diffuser(VAE,use_flux=True)
|
||||
dreamo_pipeline = DreamOPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"), vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
|
||||
elif cf_model.get("use_unet"):
|
||||
|
||||
if quantize_mode=="fp8":
|
||||
transformer = FluxTransformer2DModel.from_single_file(cf_model.get("model_path"), config=os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"),
|
||||
torch_dtype=torch.bfloat16)
|
||||
else:
|
||||
from safetensors.torch import load_file
|
||||
t_state_dict=load_file(cf_model.get("model_path"))
|
||||
unet_config = FluxTransformer2DModel.load_config(os.path.join(cur_path,"config/FLUX.1-dev/transformer/config.json"))
|
||||
transformer = FluxTransformer2DModel.from_config(unet_config).to(torch.bfloat16)
|
||||
transformer.load_state_dict(t_state_dict, strict=False)
|
||||
del t_state_dict
|
||||
gc_cleanup()
|
||||
vae=convert_cfvae2diffuser(VAE,use_flux=True)
|
||||
dreamo_pipeline = DreamOPipeline.from_pretrained(os.path.join(cur_path,"config/FLUX.1-dev"), vae=vae,transformer=transformer,text_encoder=None,text_encoder_2=None,torch_dtype=torch.bfloat16)
|
||||
else:
|
||||
print(f"use {quantize_mode} quantization")
|
||||
if quantize_mode=="fp8":
|
||||
transformer = FluxTransformer2DModel.from_pretrained(
|
||||
cf_model.get("extra_repo"),
|
||||
subfolder="transformer",
|
||||
quantization_config=DiffusersBitsAndBytesConfig(load_in_8bit=True,),
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
elif quantize_mode=="nf4": #nf4
|
||||
transformer = FluxTransformer2DModel.from_pretrained(
|
||||
cf_model.get("extra_repo"),
|
||||
subfolder="transformer",
|
||||
quantization_config=DiffusersBitsAndBytesConfig(
|
||||
load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
|
||||
),
|
||||
torch_dtype=torch.bfloat16,)
|
||||
else:
|
||||
transformer = FluxTransformer2DModel.from_pretrained(
|
||||
cf_model.get("extra_repo"),
|
||||
subfolder="transformer",
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
dreamo_pipeline = DreamOPipeline.from_pretrained(cf_model.get("extra_repo"), transformer=transformer,torch_dtype=torch.bfloat16)
|
||||
|
||||
dreamo_pipeline.load_dreamo_model(device,dreamo_lora_path,cfg_distill_path,Turbo_path, use_turbo=True)
|
||||
dreamo_pipeline.enable_model_cpu_offload()
|
||||
return dreamo_pipeline
|
||||
|
||||
@torch.no_grad()
|
||||
def get_align_face(img,face_helper):
|
||||
from .DreamO.dreamo.utils import img2tensor,tensor2img
|
||||
from torchvision.transforms.functional import normalize
|
||||
# the face preprocessing code is same as PuLID
|
||||
face_helper.clean_all()
|
||||
image_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||||
face_helper.read_image(image_bgr)
|
||||
face_helper.get_face_landmarks_5(only_center_face=True)
|
||||
face_helper.align_warp_face()
|
||||
if len(face_helper.cropped_faces) == 0:
|
||||
return None
|
||||
align_face = face_helper.cropped_faces[0]
|
||||
|
||||
input = img2tensor(align_face, bgr2rgb=True).unsqueeze(0) / 255.0
|
||||
input = input.to(torch.device("cuda"))
|
||||
parsing_out = face_helper.face_parse(normalize(input, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))[0]
|
||||
parsing_out = parsing_out.argmax(dim=1, keepdim=True)
|
||||
bg_label = [0, 16, 18, 7, 8, 9, 14, 15]
|
||||
bg = sum(parsing_out == i for i in bg_label).bool()
|
||||
white_image = torch.ones_like(input)
|
||||
# only keep the face features
|
||||
face_features_image = torch.where(bg, white_image, input)
|
||||
face_features_image = tensor2img(face_features_image, rgb2bgr=False)
|
||||
|
||||
return face_features_image
|
||||
|
||||
def Dreamo_image_encoder(BEN2_path,ref_image1,ref_image2,ref_task1,ref_task2,ref_res):
|
||||
from .DreamO.dreamo.utils import img2tensor, resize_numpy_image_area
|
||||
ref_conds = []
|
||||
#debug_images = []
|
||||
|
||||
ref_images = [ref_image1, ref_image2]
|
||||
ref_tasks = [ref_task1, ref_task2]
|
||||
|
||||
for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
|
||||
if ref_image is not None:
|
||||
if ref_task == "id":
|
||||
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
|
||||
face_helper = FaceRestoreHelper(
|
||||
upscale_factor=1,
|
||||
face_size=512,
|
||||
crop_ratio=(1, 1),
|
||||
det_model='retinaface_resnet50',
|
||||
save_ext='png',
|
||||
device=device,
|
||||
)
|
||||
ref_image = resize_numpy_image_long(ref_image, 1024)
|
||||
ref_image = get_align_face(ref_image,face_helper)
|
||||
elif ref_task != "style": # ip
|
||||
from .DreamO.tools import BEN2
|
||||
bg_rm_model = BEN2.BEN_Base().to(device).eval()
|
||||
#hf_hub_download(repo_id='PramaLLC/BEN2', filename='BEN2_Base.pth', local_dir='models')
|
||||
bg_rm_model.loadcheckpoints(BEN2_path)
|
||||
ref_image = bg_rm_model.inference(ref_image) #NEED TO CHECK
|
||||
if ref_task != "id":
|
||||
ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
|
||||
|
||||
#debug_images.append(ref_image)
|
||||
ref_image = img2tensor(ref_image, bgr2rgb=False).unsqueeze(0) / 255.0
|
||||
ref_image = 2 * ref_image - 1.0
|
||||
ref_conds.append(
|
||||
{
|
||||
'img': ref_image,
|
||||
'task': ref_task,
|
||||
'idx': idx + 1,
|
||||
}
|
||||
)
|
||||
return ref_conds
|
||||
Reference in New Issue
Block a user