| |
| import os |
| import torch |
| import torch._dynamo |
| import gc |
|
|
| torch._dynamo.config.suppress_errors = True |
| os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True" |
| os.environ["TOKENIZERS_PARALLELISM"] = "True" |
|
|
|
|
| from huggingface_hub.constants import HF_HUB_CACHE |
| from torch import Generator |
| from diffusers import FluxTransformer2DModel, DiffusionPipeline |
|
|
| from PIL.Image import Image |
| from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny |
| from pipelines.models import TextToImageRequest |
| from optimum.quanto import requantize |
| import json |
| import transformers |
| from functools import wraps |
|
|
|
|
|
|
| torch._dynamo.config.suppress_errors = True |
| os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True" |
| os.environ["TOKENIZERS_PARALLELISM"] = "True" |
|
|
| MAIN_ID = "RichardWilliam/FullyFLUXSCH" |
| REV = "c5f4f70c6cb9228a9c258799aadc660dde417af6" |
| Pipeline = None |
| apply_quanto=1 |
|
|
| def to_hell(): |
| gc.collect() |
| torch.cuda.empty_cache() |
| torch.cuda.reset_max_memory_allocated() |
| torch.cuda.reset_peak_memory_stats() |
|
|
| def error_handler(func): |
| @wraps(func) |
| def wrapper(*args, **kwargs): |
| try: |
| return func(*args, **kwargs) |
| except Exception as e: |
| print(f"Error in {func.__name__}: {str(e)}") |
| return None |
| return wrapper |
|
|
| @error_handler |
| def load_quanto_text_encoder_2(text_repo_path): |
| with open("quantization_map.json", "r") as f: |
| quantization_map = json.load(f) |
| with open(os.path.join(text_repo_path, "config.json"), "r") as f: |
| t5_config = transformers.T5Config(**json.load(f)) |
| with torch.device("meta"): |
| text_encoder_2 = transformers.T5EncoderModel(t5_config).to(torch.bfloat16) |
| state_dict = None |
| requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cuda")) |
| return text_encoder_2 |
|
|
|
|
| def load_pipeline() -> Pipeline: |
|
|
| main_path = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_Transfomer/snapshots/6860c51af40329808f270e159a0d018559a1204f") |
| origin_trans = FluxTransformer2DModel.from_pretrained(main_path, |
| torch_dtype=torch.bfloat16, |
| use_safetensors=False).to(memory_format=torch.channels_last) |
| |
| transformer = origin_trans |
|
|
| pipeline = DiffusionPipeline.from_pretrained(MAIN_ID, |
| revision=REV, |
| transformer=transformer, |
| torch_dtype=torch.bfloat16) |
| pipeline.to("cuda") |
|
|
|
|
| text_encoder_v2 = load_quanto_text_encoder_2(text_repo_path=None) |
|
|
| if text_encoder_v2==None: |
| print("Something wrong") |
| else: |
| pipeline.text_encoder_2 = text_encoder_v2 |
|
|
| for __ in range(3): |
| pipeline(prompt="I am the worst", |
| 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: |
|
|
| to_hell() |
|
|
| 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] |
