|
|
|
|
|
import os |
|
|
import torch |
|
|
import torch._dynamo |
|
|
import gc |
|
|
import json |
|
|
import transformers |
|
|
|
|
|
from huggingface_hub.constants import HF_HUB_CACHE |
|
|
from transformers import T5EncoderModel |
|
|
import diffusers |
|
|
from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only |
|
|
from torch import Generator |
|
|
from diffusers import FluxTransformer2DModel, DiffusionPipeline |
|
|
|
|
|
from PIL.Image import Image |
|
|
from diffusers import AutoencoderTiny |
|
|
from pipelines.models import TextToImageRequest |
|
|
from optimum.quanto import requantize as optimum_quant |
|
|
|
|
|
try: |
|
|
from huggingface_hub import hf_hub_download |
|
|
except: |
|
|
pass |
|
|
|
|
|
|
|
|
torch._dynamo.config.suppress_errors = True |
|
|
os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True" |
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "True" |
|
|
|
|
|
ckpt_main = "black-forest-labs/FLUX.1-schnell" |
|
|
revision_main = "741f7c3ce8b383c54771c7003378a50191e9efe9" |
|
|
Pipeline = None |
|
|
apply_transformer_tag = 1 |
|
|
|
|
|
import torch |
|
|
import gc |
|
|
import os |
|
|
import json |
|
|
import transformers |
|
|
|
|
|
|
|
|
def convert_transformer_to_int8(repo_path): |
|
|
with open("transformer_int8.json", "r") as f: |
|
|
quantization_map = json.load(f) |
|
|
with torch.device("meta"): |
|
|
transformer_config_path = os.path.join(repo_path, "config.json") |
|
|
transformer = diffusers.FluxTransformer2DModel.from_config(transformer_config_path).to(torch.bfloat16) |
|
|
|
|
|
state_dict = hf_hub_download(repo_path, "diffusion_pytorch_models.safetensors") |
|
|
|
|
|
optimum_quant(transformer, state_dict, quantization_map, device=torch.device("cuda")) |
|
|
return transformer |
|
|
|
|
|
|
|
|
def load_pipeline() -> Pipeline: |
|
|
|
|
|
original_vae = AutoencoderTiny.from_pretrained("RichardWilliam/XULF_Vae", |
|
|
revision="3ee225c539465c27adadec45c6e8af50a7397b7d", |
|
|
torch_dtype=torch.bfloat16) |
|
|
|
|
|
|
|
|
text_encoder_2 = T5EncoderModel.from_pretrained("RichardWilliam/XULF_T5_bf16", |
|
|
revision = "63a3d9ef7b586655600ac9bd4e4747d038237761", |
|
|
torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last) |
|
|
|
|
|
trans_path = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_Transfomer/snapshots/6860c51af40329808f270e159a0d018559a1204f") |
|
|
pre_quanted_trans = FluxTransformer2DModel.from_pretrained(trans_path, |
|
|
torch_dtype=torch.bfloat16, |
|
|
use_safetensors=False).to(memory_format=torch.channels_last) |
|
|
transformer = pre_quanted_trans |
|
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(ckpt_main, |
|
|
revision=revision_main, |
|
|
vae=original_vae, |
|
|
transformer=transformer, |
|
|
text_encoder_2=text_encoder_2, |
|
|
torch_dtype=torch.bfloat16) |
|
|
pipeline.to("cuda") |
|
|
try: |
|
|
pipeline.enable_int8() |
|
|
pipeline.transformer = convert_transformer_to_int8(trans_path) |
|
|
except: |
|
|
print("Use origin pipeline") |
|
|
|
|
|
for warm_up_prompt in range(3): |
|
|
pipeline(prompt="puffer, cutie, buttinsky, prototrophic, betulinamaric, quintet, tunesome, decaspermous", |
|
|
width=1024, |
|
|
height=1024, |
|
|
guidance_scale=0.0, |
|
|
num_inference_steps=4, |
|
|
max_sequence_length=256) |
|
|
return pipeline |
|
|
|
|
|
@torch.no_grad() |
|
|
def infer(request: TextToImageRequest, pipeline: Pipeline) -> Image: |
|
|
|
|
|
gc.collect() |
|
|
torch.cuda.empty_cache() |
|
|
|
|
|
generator = Generator(pipeline.device).manual_seed(request.seed) |
|
|
|
|
|
return pipeline( |
|
|
request.prompt, |
|
|
generator=generator, |
|
|
guidance_scale=0.0, |
|
|
num_inference_steps=4, |
|
|
max_sequence_length=256, |
|
|
height=request.height, |
|
|
width=request.width, |
|
|
).images[0] |
