import os import torch import numpy as np from PIL import Image from einops import rearrange, repeat from typing import Optional import math import torch.nn as nn import gc from diffusers.pipelines.flux.modeling_flux import ReduxImageEncoder from .src.flux.sampling import denoise_lucidflux, get_noise, get_schedule, unpack from .src.flux.util import load_flow_model, load_single_condition_branch, load_safetensors from .src.flux.swinir import SwinIR def get_memory_status(): """Get current GPU memory usage""" if torch.cuda.is_available(): allocated = torch.cuda.memory_allocated() / 1024**3 reserved = torch.cuda.memory_reserved() / 1024**3 free = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved()) / 1024**3 return { "allocated": allocated, "reserved": reserved, "free": free, "total": torch.cuda.get_device_properties(0).total_memory / 1024**3 } return None def aggressive_cleanup(): gc.collect() gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() torch.cuda.synchronize() class ComfyUIFluxWrapper(nn.Module): def __init__(self, comfyui_model): super().__init__() self.model = comfyui_model def forward(self, img, img_ids, txt, txt_ids, y, timesteps, guidance, block_controlnet_hidden_states=None): try: b, seq_len, latent_dim = img.shape h = int((img_ids[0, :, 1].max().item() + 1)) w = int((img_ids[0, :, 2].max().item() + 1)) c = latent_dim // 4 # 4 = ph * pw (2 * 2) img_unpacked = rearrange(img, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h, w=w, ph=2, pw=2, c=c) control_dict = None if block_controlnet_hidden_states is not None and len(block_controlnet_hidden_states) > 0: control_dict = { "input": [], "middle": [], "output": block_controlnet_hidden_states } output_unpacked = self.model( x=img_unpacked, timestep=timesteps, context=txt, y=y, guidance=guidance, txt_ids=txt_ids, img_ids=img_ids, control=control_dict, transformer_options={} ) output_packed = rearrange(output_unpacked, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) return output_packed except Exception as e: print(f"ComfyUIFluxWrapper error: {e}") raise def to(self, *args, **kwargs): """Override to() to move the wrapped model""" self.model = self.model.to(*args, **kwargs) return super().to(*args, **kwargs) def cpu(self): """Override cpu() to move the wrapped model""" self.model = self.model.cpu() return super().cpu() def cuda(self, device=None): """Override cuda() to move the wrapped model""" self.model = self.model.cuda(device) return super().cuda(device) def state_dict(self, *args, **kwargs): return self.model.state_dict(*args, **kwargs) def load_state_dict(self, state_dict, *args, **kwargs): return self.model.load_state_dict(state_dict, *args, **kwargs) def print_memory_status(prefix=""): mem = get_memory_status() if mem: print(f"{prefix}VRAM: {mem['allocated']:.2f}GB allocated | {mem['reserved']:.2f}GB reserved | {mem['free']:.2f}GB free | {mem['total']:.2f}GB total") def get_timestep_embedding( timesteps: torch.Tensor, embedding_dim: int, flip_sin_to_cos: bool = False, downscale_freq_shift: float = 1, scale: float = 1, max_period: int = 10000, ): assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" half_dim = embedding_dim // 2 exponent = -math.log(max_period) * torch.arange( start=0, end=half_dim, dtype=torch.float32, device=timesteps.device ) exponent = exponent / (half_dim - downscale_freq_shift) emb = torch.exp(exponent) emb = timesteps[:, None].float() * emb[None, :] emb = scale * emb emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) if flip_sin_to_cos: emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) if embedding_dim % 2 == 1: emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) return emb class Timesteps(nn.Module): def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): super().__init__() self.num_channels = num_channels self.flip_sin_to_cos = flip_sin_to_cos self.downscale_freq_shift = downscale_freq_shift self.scale = scale def forward(self, timesteps): t_emb = get_timestep_embedding( timesteps, self.num_channels, flip_sin_to_cos=self.flip_sin_to_cos, downscale_freq_shift=self.downscale_freq_shift, scale=self.scale, ) return t_emb ACT2CLS = { "swish": nn.SiLU, "silu": nn.SiLU, "mish": nn.Mish, "gelu": nn.GELU, "relu": nn.ReLU, } def get_activation(act_fn: str) -> nn.Module: act_fn = act_fn.lower() if act_fn in ACT2CLS: return ACT2CLS[act_fn]() else: raise ValueError(f"activation function {act_fn} not found in ACT2FN mapping {list(ACT2CLS.keys())}") class TimestepEmbedding(nn.Module): def __init__( self, in_channels: int, time_embed_dim: int, act_fn: str = "silu", out_dim: int = None, post_act_fn: Optional[str] = None, cond_proj_dim=None, sample_proj_bias=True, ): super().__init__() self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) if cond_proj_dim is not None: self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) else: self.cond_proj = None self.act = get_activation(act_fn) if out_dim is not None: time_embed_dim_out = out_dim else: time_embed_dim_out = time_embed_dim self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) if post_act_fn is None: self.post_act = None else: self.post_act = get_activation(post_act_fn) def forward(self, sample, condition=None): if condition is not None: sample = sample + self.cond_proj(condition) sample = self.linear_1(sample) if self.act is not None: sample = self.act(sample) sample = self.linear_2(sample) if self.post_act is not None: sample = self.post_act(sample) return sample class Modulation(nn.Module): def __init__(self, dim, bias=True): super().__init__() self.silu = nn.SiLU() self.linear = nn.Linear(dim, 2 * dim, bias=bias) self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=dim) self.control_index_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=dim) def forward(self, x, timestep, control_index): timesteps_proj = self.time_proj(timestep * 1000) timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=x.dtype)) # (N, D) if control_index.dim() == 0: control_index = control_index.repeat(x.shape[0]) elif control_index.dim() == 1 and control_index.shape[0] != x.shape[0]: control_index = control_index.expand(x.shape[0]) control_index = control_index.to(device=x.device, dtype=x.dtype) control_index_proj = self.time_proj(control_index) control_index_emb = self.control_index_embedder(control_index_proj.to(dtype=x.dtype)) # (N, D) timesteps_emb = timesteps_emb + control_index_emb emb = self.linear(self.silu(timesteps_emb)) shift_msa, scale_msa = emb.chunk(2, dim=1) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x class DualConditionBranch(nn.Module): def __init__(self, condition_branch_lq: nn.Module, condition_branch_ldr: nn.Module, modulation_lq: nn.Module, modulation_ldr: nn.Module): super().__init__() self.lq = condition_branch_lq self.ldr = condition_branch_ldr self.modulation_lq = modulation_lq self.modulation_ldr = modulation_ldr def forward( self, *, img, img_ids, condition_cond_lq, txt, txt_ids, y, timesteps, guidance, condition_cond_ldr=None, ): out_lq = self.lq( img=img, img_ids=img_ids, controlnet_cond=condition_cond_lq, txt=txt, txt_ids=txt_ids, y=y, timesteps=timesteps, guidance=guidance, ) out_ldr = self.ldr( img=img, img_ids=img_ids, controlnet_cond=condition_cond_ldr, txt=txt, txt_ids=txt_ids, y=y, timesteps=timesteps, guidance=guidance, ) out = [] num_blocks = 19 for i in range(num_blocks // 2 + 1): for control_index, (lq, ldr) in enumerate(zip(out_lq, out_ldr)): control_index = torch.tensor(control_index, device=timesteps.device, dtype=timesteps.dtype) lq = self.modulation_lq(lq, timesteps, i * 2 + control_index) if len(out) == num_blocks: break ldr = self.modulation_ldr(ldr, timesteps, i * 2 + control_index) out.append(lq + ldr) return out def preprocess_lq_image(image_path: str, width: int = 512, height: int = 512): image = Image.open(image_path).convert('RGB') image = image.resize((width, height)) return image def load_redux_image_encoder(device: torch.device, dtype: torch.dtype, redux_state_dict: str): redux_image_encoder = ReduxImageEncoder() redux_image_encoder.load_state_dict(redux_state_dict, strict=False) redux_image_encoder.eval() redux_image_encoder.to(device).to(dtype=dtype) return redux_image_encoder def load_lucidflux_model(args,ckpt_path,cf_model,torch_device,offload=False): name =args.name is_schnell = name == "flux-schnell" model = load_flow_model(name, ckpt_path, cf_model) if type(model).__module__.startswith('comfy.'): model = ComfyUIFluxWrapper(model) if not offload: model = model.to(torch_device) device_for_loading = 'cpu' if offload else torch_device condition_lq=load_single_condition_branch(name, device_for_loading).to(torch.bfloat16) if '.safetensors' in args.checkpoint: checkpoint = load_safetensors(args.checkpoint) else: checkpoint = torch.load(args.checkpoint,weights_only=False, map_location='cpu') condition_lq.load_state_dict(checkpoint["condition_lq"], strict=False) if not offload: condition_lq = condition_lq.to(torch_device) condition_ldr = load_single_condition_branch(name, device_for_loading).to(torch.bfloat16) condition_ldr.load_state_dict(checkpoint["condition_ldr"], strict=False) if not offload: condition_ldr = condition_ldr.to(torch_device) modulation_lq = Modulation(dim=3072).to(torch.bfloat16) modulation_lq.load_state_dict(checkpoint["modulation_lq"], strict=False) modulation_ldr = Modulation(dim=3072).to(torch.bfloat16) modulation_ldr.load_state_dict(checkpoint["modulation_ldr"], strict=False) dual_condition_branch = DualConditionBranch( condition_lq, condition_ldr, modulation_lq=modulation_lq, modulation_ldr=modulation_ldr, ) if not offload: dual_condition_branch = dual_condition_branch.to(torch_device) state_dict=checkpoint["connector"] del checkpoint gc.collect() if torch_device.type == 'cuda': torch.cuda.empty_cache() pipe={"model":model,"dual_condition_branch":dual_condition_branch,"is_schnell":is_schnell,"offload":offload} return pipe,state_dict def preprocess_data(state_dict,swinir_path,siglip_model,input_pli_list, inp_cond,torch_device): swinir = SwinIR( img_size=64, patch_size=1, in_chans=3, embed_dim=180, depths=[6, 6, 6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6, 6, 6], window_size=8, mlp_ratio=2, sf=8, img_range=1.0, upsampler="nearest+conv", resi_connection="1conv", unshuffle=True, unshuffle_scale=8, ) ckpt_obj = torch.load(swinir_path, weights_only=False,map_location="cpu") state = ckpt_obj.get("state_dict", ckpt_obj) new_state = {k.replace("module.", ""): v for k, v in state.items()} swinir.load_state_dict(new_state, strict=False) swinir.eval() del state for p in swinir.parameters(): p.requires_grad_(False) swinir = swinir.to(torch_device) dtype = torch.bfloat16 if torch_device.type == 'cuda' else torch.float32 redux_image_encoder = load_redux_image_encoder(torch_device, dtype, state_dict) del state_dict data_list=[] for lq_processed in input_pli_list: condition_cond = torch.from_numpy((np.array(lq_processed) / 127.5) - 1) condition_cond = condition_cond.permute(2, 0, 1).unsqueeze(0).to(torch.bfloat16).to(torch_device) condition_cond_ldr = None with torch.no_grad(): ci_01 = torch.clamp((condition_cond.float() + 1.0) / 2.0, 0.0, 1.0) ci_pre = swinir(ci_01).float().clamp(0.0, 1.0) #(1,3,H,W) or 3 h w ci_pre_origin=ci_pre condition_cond_ldr = (ci_pre * 2.0 - 1.0).to(torch.bfloat16) if ci_pre.ndim == 3: ci_pre = ci_pre.unsqueeze(0) _,_,height,width=ci_pre.shape ci_pre=ci_pre.permute(0, 2, 3, 1) # to comfy siglip_image_pre_fts=siglip_model.encode_image(ci_pre)["last_hidden_state"].to(device=torch_device,dtype=torch.bfloat16) enc_dtype = redux_image_encoder.redux_up.weight.dtype image_embeds = redux_image_encoder( siglip_image_pre_fts.to(device=torch_device, dtype=enc_dtype) )["image_embeds"] txt = inp_cond["txt"].to(device=torch_device, dtype=torch.bfloat16) txt_ids = inp_cond["txt_ids"].to(device=torch_device, dtype=torch.bfloat16) siglip_txt = torch.cat([txt, image_embeds.to(dtype=torch.bfloat16)], dim=1).to(device=torch_device, dtype=torch.bfloat16) B, L, C = txt_ids.shape extra_ids = torch.zeros((B, 1024, C), device=txt_ids.device, dtype=torch.bfloat16) siglip_txt_ids = torch.cat([txt_ids, extra_ids], dim=1).to(device=torch_device,dtype=torch.bfloat16) data={"siglip_txt": siglip_txt, "siglip_txt_ids": siglip_txt_ids,"inp_cond":inp_cond,"txt":txt,"txt_ids":txt_ids,"size":(height, width), "condition_cond":condition_cond, "condition_cond_ldr": condition_cond_ldr,"ci_pre_origin":ci_pre_origin} data_list.append(data) return data_list def get_cond(positive,height,width,device,bs=1): h=2 * math.ceil(height / 16) w=2 * math.ceil(width / 16) img_ids = torch.zeros(h // 2, w // 2, 3) img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None] img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :] img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) if positive is None: txt = torch.zeros(bs, 256, 4096) txt_ids = torch.zeros(bs, 256, 3) vec = torch.zeros(bs, 768) else: txt = positive[0][0] if txt.shape[0] == 1 and bs > 1: txt = repeat(txt, "1 ... -> bs ...", bs=bs) txt_ids = torch.zeros(bs, txt.shape[1], 3) vec = positive[0][1].get("pooled_output") if vec.shape[0] == 1 and bs > 1: vec = repeat(vec, "1 ... -> bs ...", bs=bs) inp_cond={ "img_ids": img_ids.to(device), "txt": txt.to(device), "txt_ids": txt_ids.to(device), "vec": vec.to(device), } return inp_cond def get_cond_from_embeddings(prompt_embeddings, height, width, device, bs=1): h = 2 * math.ceil(height / 16) w = 2 * math.ceil(width / 16) img_ids = torch.zeros(h // 2, w // 2, 3) img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None] img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :] img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs) if isinstance(prompt_embeddings, dict): txt = prompt_embeddings.get('txt', prompt_embeddings.get('text_embeddings')) vec = prompt_embeddings.get('vec', prompt_embeddings.get('pooled_embeddings')) else: txt = prompt_embeddings vec = torch.zeros(bs, 768) if txt.shape[0] == 1 and bs > 1: txt = repeat(txt, "1 ... -> bs ...", bs=bs) if vec.shape[0] == 1 and bs > 1: vec = repeat(vec, "1 ... -> bs ...", bs=bs) txt_ids = torch.zeros(bs, txt.shape[1], 3) inp_cond = { "img_ids": img_ids.to(device), "txt": txt.to(device), "txt_ids": txt_ids.to(device), "vec": vec.to(device), } return inp_cond def preprocess_data_cached(swinir_cache, redux_cache, state_dict, swinir_path, siglip_model, input_pli_list, inp_cond, torch_device, offload=False): dtype = torch.bfloat16 if torch_device.type == 'cuda' else torch.float32 if swinir_path not in swinir_cache: swinir = SwinIR( img_size=64, patch_size=1, in_chans=3, embed_dim=180, depths=[6, 6, 6, 6, 6, 6, 6, 6], num_heads=[6, 6, 6, 6, 6, 6, 6, 6], window_size=8, mlp_ratio=2, sf=8, img_range=1.0, upsampler="nearest+conv", resi_connection="1conv", unshuffle=True, unshuffle_scale=8, ) ckpt_obj = torch.load(swinir_path, weights_only=False, map_location="cpu") state = ckpt_obj.get("state_dict", ckpt_obj) new_state = {k.replace("module.", ""): v for k, v in state.items()} swinir.load_state_dict(new_state, strict=False) swinir.eval() for p in swinir.parameters(): p.requires_grad_(False) del state, ckpt_obj aggressive_cleanup() swinir_cache[swinir_path] = swinir else: swinir = swinir_cache[swinir_path] swinir = swinir.to('cpu' if offload else torch_device) swinir_results = [] for idx, lq_processed in enumerate(input_pli_list): condition_cond = torch.from_numpy((np.array(lq_processed) / 127.5) - 1) condition_cond = condition_cond.permute(2, 0, 1).unsqueeze(0).to(dtype).to(torch_device) with torch.no_grad(): ci_01 = torch.clamp((condition_cond.float() + 1.0) / 2.0, 0.0, 1.0) if offload: ci_01_input = ci_01.cpu() ci_pre = swinir(ci_01_input).float().clamp(0.0, 1.0) ci_pre = ci_pre.to(torch_device) del ci_01_input else: ci_pre = swinir(ci_01).float().clamp(0.0, 1.0) del ci_01 ci_pre_origin = ci_pre.clone() condition_cond_ldr = (ci_pre * 2.0 - 1.0).to(dtype) swinir_results.append({ "condition_cond": condition_cond, "ci_pre": ci_pre, "ci_pre_origin": ci_pre_origin, "condition_cond_ldr": condition_cond_ldr }) if idx % 1 == 0: aggressive_cleanup() swinir_cache[swinir_path] = swinir.cpu() del swinir aggressive_cleanup() siglip_results = [] for idx, result in enumerate(swinir_results): ci_pre = result["ci_pre"] with torch.no_grad(): if ci_pre.ndim == 3: ci_pre = ci_pre.unsqueeze(0) _, _, height, width = ci_pre.shape ci_pre_perm = ci_pre.permute(0, 2, 3, 1) siglip_image_pre_fts = siglip_model.encode_image(ci_pre_perm)["last_hidden_state"] siglip_image_pre_fts = siglip_image_pre_fts.to(device=torch_device, dtype=dtype) del ci_pre_perm siglip_results.append({ "siglip_features": siglip_image_pre_fts, "size": (height, width) }) if idx % 1 == 0: aggressive_cleanup() redux_key = str(id(state_dict)) if redux_key not in redux_cache: redux_image_encoder = load_redux_image_encoder(torch_device, dtype, state_dict) redux_cache[redux_key] = redux_image_encoder.cpu() else: redux_image_encoder = redux_cache[redux_key] redux_image_encoder = redux_image_encoder.to(torch_device).to(dtype) data_list = [] for idx, (swinir_res, siglip_res) in enumerate(zip(swinir_results, siglip_results)): with torch.no_grad(): siglip_image_pre_fts = siglip_res["siglip_features"] enc_dtype = redux_image_encoder.redux_up.weight.dtype if siglip_image_pre_fts.dtype != enc_dtype: siglip_image_pre_fts = siglip_image_pre_fts.to(enc_dtype) image_embeds = redux_image_encoder(siglip_image_pre_fts)["image_embeds"] txt = inp_cond["txt"].to(device=torch_device, dtype=dtype) txt_ids = inp_cond["txt_ids"].to(device=torch_device, dtype=dtype) siglip_txt = torch.cat([txt, image_embeds.to(dtype)], dim=1) B, L, C = txt_ids.shape extra_ids = torch.zeros((B, 1024, C), device=txt_ids.device, dtype=dtype) siglip_txt_ids = torch.cat([txt_ids, extra_ids], dim=1) data = { "siglip_txt": siglip_txt, "siglip_txt_ids": siglip_txt_ids, "inp_cond": inp_cond, "txt": txt, "txt_ids": txt_ids, "size": siglip_res["size"], "condition_cond": swinir_res["condition_cond"], "condition_cond_ldr": swinir_res["condition_cond_ldr"], "ci_pre_origin": swinir_res["ci_pre_origin"] } del siglip_image_pre_fts, image_embeds data_list.append(data) aggressive_cleanup() del swinir_results, siglip_results redux_cache[redux_key] = redux_image_encoder.cpu() del redux_image_encoder aggressive_cleanup() return data_list def lucidflux_inference(model,dual_condition_branch,input_data,guidance,num_steps,seed,torch_device,is_schnell=False,offload=False,progress_callback=None): lat_list = [] for idx, data in enumerate(input_data): with torch.no_grad(): height, width=data.get("size") torch.manual_seed(seed) x = get_noise( 1, height, width, device=torch_device, dtype=torch.bfloat16, seed=seed ) bs, c, h, w = x.shape img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) if img.shape[0] == 1 and bs > 1: img = repeat(img, "1 ... -> bs ...", bs=bs) timesteps = get_schedule( num_steps, (width // 8) * (height // 8) // (16 * 16), shift=(not is_schnell), ) # Move models to device if offloading if offload: model = model.to(torch_device) dual_condition_branch.to(torch_device) x = denoise_lucidflux( model, dual_condition_model=dual_condition_branch, img=img, img_ids=data.get("inp_cond")["img_ids"], txt=data.get("txt"), txt_ids=data.get("txt_ids"), siglip_txt=data.get("siglip_txt"), siglip_txt_ids=data.get("siglip_txt_ids"), vec=data.get("inp_cond")["vec"], timesteps=timesteps, guidance=guidance, condition_cond_lq=data.get("condition_cond"), condition_cond_ldr=data.get("condition_cond_ldr"), progress_callback=progress_callback, ) # Offload models back to CPU if offload: model = model.cpu() dual_condition_branch.cpu() aggressive_cleanup() x = unpack(x.float(), height, width) x=(x/0.3611)+0.1159 #mean lat_list.append(x) del img, timesteps aggressive_cleanup() return lat_list