File size: 3,653 Bytes
a06bf5a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from huggingface_hub.constants import HF_HUB_CACHE
from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
import torch
import torch._dynamo
import gc
import os
from torch import Generator
from diffusers import FluxTransformer2DModel, DiffusionPipeline
from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
from PIL.Image import Image
from pipelines.models import TextToImageRequest
# Add env optimize
os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "True"
torch._dynamo.config.suppress_errors = True
Pipeline = None
CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
class NormQuant:
def __init__(self, model, noise_level=0.05):
self.model = model
self.noise_level = noise_level
def apply(self):
for name, param in self.model.named_parameters():
if param.requires_grad:
with torch.no_grad():
noise = torch.randn_like(param.data) * self.noise_level
param.data = torch.floor(param.data + noises)
for buffer_name, buffer in self.model.named_buffers():
with torch.no_grad():
buffer.add_(torch.full_like(buffer, 0.01))
return self.model
def load_pipeline() -> Pipeline:
vae = AutoencoderTiny.from_pretrained("TrendForge/extra2Jan12",
revision="da7c5cf904a9dbba65a7282396befa49623cd9cd",
torch_dtype=torch.bfloat16)
base_text_encoder_2 = T5EncoderModel.from_pretrained("TrendForge/extra1Jan11",
revision = "c76831ddf0852be22835f79dc5c1fbacb1ccda9e",
torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
# Apply to text_encoder_2
try:
text_encoder_2 = NormQuant(base_text_encoder_2, noise_level=0.03).apply()
except:
text_encoder_2 = base_text_encoder_2
path = os.path.join(HF_HUB_CACHE, "models--TrendForge--extra0Jan10/snapshots/d3ded25a77fdef06de4059d94b080a34da6e7a82")
base_transformer = FluxTransformer2DModel.from_pretrained(path,
torch_dtype=torch.bfloat16,
use_safetensors=False).to(memory_format=torch.channels_last)
# Apply to transformer
try:
transformer = NormQuant(base_transformer, noise_level=0.03).apply()
except:
transformer = base_transformer
pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT,
revision=REVISION,
vae=vae,
transformer=transformer,
text_encoder_2=text_encoder_2,
torch_dtype=torch.bfloat16)
pipeline.to("cuda")
for _ in range(3):
pipeline(prompt="freezable, catacorolla, gaiassa, unenkindled, grubs, solidiform",
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:
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] |