Spaces:
Sleeping
Sleeping
import sys | |
import os | |
from pathlib import Path | |
import gc | |
# Add the StableCascade and CSD directories to the Python path | |
app_dir = Path(__file__).parent | |
sys.path.extend([ | |
str(app_dir), | |
str(app_dir / "third_party" / "StableCascade"), | |
str(app_dir / "third_party" / "CSD") | |
]) | |
import yaml | |
import torch | |
from tqdm import tqdm | |
from accelerate.utils import set_module_tensor_to_device | |
import torch.nn.functional as F | |
import torchvision.transforms as T | |
from lang_sam import LangSAM | |
from inference.utils import * | |
from core.utils import load_or_fail | |
from train import WurstCoreC, WurstCoreB | |
from gdf_rbm import RBM | |
from stage_c_rbm import StageCRBM | |
from utils import WurstCoreCRBM | |
from gdf.schedulers import CosineSchedule | |
from gdf import VPScaler, CosineTNoiseCond, DDPMSampler, P2LossWeight, AdaptiveLossWeight | |
from gdf.targets import EpsilonTarget | |
import PIL | |
# Device configuration | |
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
print(device) | |
# Flag for low VRAM usage | |
low_vram = True | |
# Function definition for low VRAM usage | |
def models_to(model, device="cpu", excepts=None): | |
""" | |
Change the device of nn.Modules within a class, skipping specified attributes. | |
""" | |
for attr_name in dir(model): | |
if attr_name.startswith('__') and attr_name.endswith('__'): | |
continue # skip special attributes | |
attr_value = getattr(model, attr_name, None) | |
if isinstance(attr_value, torch.nn.Module): | |
if excepts and attr_name in excepts: | |
print(f"Except '{attr_name}'") | |
continue | |
print(f"Change device of '{attr_name}' to {device}") | |
attr_value.to(device) | |
torch.cuda.empty_cache() | |
# Stage C model configuration | |
config_file = 'third_party/StableCascade/configs/inference/stage_c_3b.yaml' | |
with open(config_file, "r", encoding="utf-8") as file: | |
loaded_config = yaml.safe_load(file) | |
core = WurstCoreCRBM(config_dict=loaded_config, device=device, training=False) | |
# Stage B model configuration | |
config_file_b = 'third_party/StableCascade/configs/inference/stage_b_3b.yaml' | |
with open(config_file_b, "r", encoding="utf-8") as file: | |
config_file_b = yaml.safe_load(file) | |
core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False) | |
# Setup extras and models for Stage C | |
extras = core.setup_extras_pre() | |
gdf_rbm = RBM( | |
schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]), | |
input_scaler=VPScaler(), target=EpsilonTarget(), | |
noise_cond=CosineTNoiseCond(), | |
loss_weight=AdaptiveLossWeight(), | |
) | |
sampling_configs = { | |
"cfg": 5, | |
"sampler": DDPMSampler(gdf_rbm), | |
"shift": 1, | |
"timesteps": 20 | |
} | |
extras = core.Extras( | |
gdf=gdf_rbm, | |
sampling_configs=sampling_configs, | |
transforms=extras.transforms, | |
effnet_preprocess=extras.effnet_preprocess, | |
clip_preprocess=extras.clip_preprocess | |
) | |
models = core.setup_models(extras) | |
models.generator.eval().requires_grad_(False) | |
# Setup extras and models for Stage B | |
extras_b = core_b.setup_extras_pre() | |
models_b = core_b.setup_models(extras_b, skip_clip=True) | |
models_b = WurstCoreB.Models( | |
**{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model} | |
) | |
models_b.generator.bfloat16().eval().requires_grad_(False) | |
# Off-load old generator (low VRAM mode) | |
if low_vram: | |
models.generator.to("cpu") | |
torch.cuda.empty_cache() | |
# Load and configure new generator | |
generator_rbm = StageCRBM() | |
for param_name, param in load_or_fail(core.config.generator_checkpoint_path).items(): | |
set_module_tensor_to_device(generator_rbm, param_name, "cpu", value=param) | |
generator_rbm = generator_rbm.to(getattr(torch, core.config.dtype)).to(device) | |
generator_rbm = core.load_model(generator_rbm, 'generator') | |
# Create models_rbm instance | |
models_rbm = core.Models( | |
effnet=models.effnet, | |
text_model=models.text_model, | |
tokenizer=models.tokenizer, | |
generator=generator_rbm, | |
previewer=models.previewer, | |
image_model=models.image_model # Add this line | |
) | |
def reset_inference_state(): | |
global models_rbm, models_b, extras, extras_b, device, core, core_b | |
# Reset sampling configurations | |
extras.sampling_configs['cfg'] = 5 | |
extras.sampling_configs['shift'] = 1 | |
extras.sampling_configs['timesteps'] = 20 | |
extras.sampling_configs['t_start'] = 1.0 | |
extras_b.sampling_configs['cfg'] = 1.1 | |
extras_b.sampling_configs['shift'] = 1 | |
extras_b.sampling_configs['timesteps'] = 10 | |
extras_b.sampling_configs['t_start'] = 1.0 | |
# Move models to CPU to free up GPU memory | |
models_to(models_rbm, device="cpu") | |
models_b.generator.to("cpu") | |
# Clear CUDA cache | |
torch.cuda.empty_cache() | |
gc.collect() | |
# Move necessary models back to the correct device | |
if low_vram: | |
models_to(models_rbm, device="cpu", excepts=["generator", "previewer"]) | |
models_rbm.generator.to(device) | |
models_rbm.previewer.to(device) | |
else: | |
models_to(models_rbm, device=device) | |
models_b.generator.to("cpu") # Keep Stage B generator on CPU for now | |
# Ensure effnet and image_model are on the correct device | |
models_rbm.effnet.to(device) | |
if models_rbm.image_model is not None: | |
models_rbm.image_model.to(device) | |
# Reset model states | |
models_rbm.generator.eval().requires_grad_(False) | |
models_b.generator.bfloat16().eval().requires_grad_(False) | |
# Clear CUDA cache again | |
torch.cuda.empty_cache() | |
gc.collect() | |
def infer(ref_style_file, style_description, caption): | |
global models_rbm, models_b | |
try: | |
caption = f"{caption} in {style_description}" | |
height=1024 | |
width=1024 | |
batch_size=1 | |
output_file='output.png' | |
stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size) | |
extras.sampling_configs['cfg'] = 4 | |
extras.sampling_configs['shift'] = 2 | |
extras.sampling_configs['timesteps'] = 20 | |
extras.sampling_configs['t_start'] = 1.0 | |
extras_b.sampling_configs['cfg'] = 1.1 | |
extras_b.sampling_configs['shift'] = 1 | |
extras_b.sampling_configs['timesteps'] = 10 | |
extras_b.sampling_configs['t_start'] = 1.0 | |
ref_style = resize_image(PIL.Image.open(ref_style_file).convert("RGB")).unsqueeze(0).expand(batch_size, -1, -1, -1).to(device) | |
batch = {'captions': [caption] * batch_size} | |
batch['style'] = ref_style | |
# Ensure models are on the correct device before inference | |
if low_vram: | |
models_to(models_rbm, device=device, excepts=["generator", "previewer"]) | |
else: | |
models_to(models_rbm, device=device) | |
models_b.generator.to(device) | |
# Ensure effnet and image_model are on the correct device | |
models_rbm.effnet.to(device) | |
if models_rbm.image_model is not None: | |
models_rbm.image_model.to(device) | |
x0_style_forward = models_rbm.effnet(extras.effnet_preprocess(ref_style)) | |
conditions = core.get_conditions(batch, models_rbm, extras, is_eval=True, is_unconditional=False, eval_image_embeds=True, eval_style=True, eval_csd=False) | |
unconditions = core.get_conditions(batch, models_rbm, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False) | |
conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False) | |
unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True) | |
if low_vram: | |
# The sampling process uses more vram, so we offload everything except two modules to the cpu. | |
models_to(models_rbm, device="cpu", excepts=["generator", "previewer"]) | |
# Stage C reverse process. | |
sampling_c = extras.gdf.sample( | |
models_rbm.generator, conditions, stage_c_latent_shape, | |
unconditions, device=device, | |
**extras.sampling_configs, | |
x0_style_forward=x0_style_forward, | |
apply_pushforward=False, tau_pushforward=8, | |
num_iter=3, eta=0.1, tau=20, eval_csd=True, | |
extras=extras, models=models_rbm, | |
lam_style=1, lam_txt_alignment=1.0, | |
use_ddim_sampler=True, | |
) | |
for (sampled_c, _, _) in tqdm(sampling_c, total=extras.sampling_configs['timesteps']): | |
sampled_c = sampled_c | |
# Stage B reverse process. | |
with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16): | |
conditions_b['effnet'] = sampled_c | |
unconditions_b['effnet'] = torch.zeros_like(sampled_c) | |
sampling_b = extras_b.gdf.sample( | |
models_b.generator, conditions_b, stage_b_latent_shape, | |
unconditions_b, device=device, **extras_b.sampling_configs, | |
) | |
for (sampled_b, _, _) in tqdm(sampling_b, total=extras_b.sampling_configs['timesteps']): | |
sampled_b = sampled_b | |
sampled = models_b.stage_a.decode(sampled_b).float() | |
sampled = torch.cat([ | |
torch.nn.functional.interpolate(ref_style.cpu(), size=(height, width)), | |
sampled.cpu(), | |
], dim=0) | |
# Remove the batch dimension and keep only the generated image | |
sampled = sampled[1] # This selects the generated image, discarding the reference style image | |
# Ensure the tensor is in [C, H, W] format | |
if sampled.dim() == 3 and sampled.shape[0] == 3: | |
sampled_image = T.ToPILImage()(sampled) # Convert tensor to PIL image | |
sampled_image.save(output_file) # Save the image as a PNG | |
else: | |
raise ValueError(f"Expected tensor of shape [3, H, W] but got {sampled.shape}") | |
return output_file # Return the path to the saved image | |
finally: | |
# Reset the state after inference, regardless of success or failure | |
reset_inference_state() | |
def reset_compo_inference_state(): | |
global models_rbm, models_b, extras, extras_b, device, core, core_b, sam_model | |
# Reset sampling configurations | |
extras.sampling_configs['cfg'] = 4 | |
extras.sampling_configs['shift'] = 2 | |
extras.sampling_configs['timesteps'] = 20 | |
extras.sampling_configs['t_start'] = 1.0 | |
extras_b.sampling_configs['cfg'] = 1.1 | |
extras_b.sampling_configs['shift'] = 1 | |
extras_b.sampling_configs['timesteps'] = 10 | |
extras_b.sampling_configs['t_start'] = 1.0 | |
# Move models to CPU to free up GPU memory | |
models_to(models_rbm, device="cpu") | |
models_b.generator.to("cpu") | |
# Clear CUDA cache | |
torch.cuda.empty_cache() | |
gc.collect() | |
# Move SAM model components to CPU if they exist | |
models_to(sam_model, device="cpu") | |
models_to(sam_model.sam, device="cpu") | |
# Clear CUDA cache | |
torch.cuda.empty_cache() | |
gc.collect() | |
# Ensure all models are in eval mode and don't require gradients | |
for model in [models_rbm.generator, models_b.generator]: | |
model.eval() | |
for param in model.parameters(): | |
param.requires_grad = False | |
# Clear CUDA cache again | |
torch.cuda.empty_cache() | |
gc.collect() | |
def infer_compo(style_description, ref_style_file, caption, ref_sub_file): | |
global models_rbm, models_b, device, sam_model | |
try: | |
caption = f"{caption} in {style_description}" | |
sam_prompt = f"{caption}" | |
use_sam_mask = False | |
# Ensure all models are on the correct device | |
models_to(models_rbm, device) | |
models_b.generator.to(device) | |
batch_size = 1 | |
height, width = 1024, 1024 | |
stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size) | |
ref_style = resize_image(PIL.Image.open(ref_style_file).convert("RGB")).unsqueeze(0).expand(batch_size, -1, -1, -1).to(device) | |
ref_images = resize_image(PIL.Image.open(ref_sub_file).convert("RGB")).unsqueeze(0).expand(batch_size, -1, -1, -1).to(device) | |
batch = {'captions': [caption] * batch_size, 'style': ref_style, 'images': ref_images} | |
x0_forward = models_rbm.effnet(extras.effnet_preprocess(ref_images)) | |
x0_style_forward = models_rbm.effnet(extras.effnet_preprocess(ref_style)) | |
## SAM Mask for sub | |
use_sam_mask = False | |
x0_preview = models_rbm.previewer(x0_forward) | |
sam_model = LangSAM() | |
# Move SAM model components to the correct device | |
models_to(sam_model, device) | |
models_to(sam_model.sam, device) | |
x0_preview_pil = T.ToPILImage()(x0_preview[0].cpu()) | |
sam_mask, boxes, phrases, logits = sam_model.predict(x0_preview_pil, sam_prompt) | |
sam_mask = sam_mask.detach().unsqueeze(dim=0).to(device) | |
conditions = core.get_conditions(batch, models_rbm, extras, is_eval=True, is_unconditional=False, eval_image_embeds=True, eval_subject_style=True, eval_csd=False) | |
unconditions = core.get_conditions(batch, models_rbm, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False, eval_subject_style=True) | |
conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False) | |
unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True) | |
if low_vram: | |
models_to(models_rbm, device="cpu", excepts=["generator", "previewer"]) | |
models_to(sam_model, device="cpu") | |
models_to(sam_model.sam, device="cpu") | |
# Stage C reverse process. | |
sampling_c = extras.gdf.sample( | |
models_rbm.generator, conditions, stage_c_latent_shape, | |
unconditions, device=device, | |
**extras.sampling_configs, | |
x0_style_forward=x0_style_forward, x0_forward=x0_forward, | |
apply_pushforward=False, tau_pushforward=5, tau_pushforward_csd=10, | |
num_iter=3, eta=1e-1, tau=20, eval_sub_csd=True, | |
extras=extras, models=models_rbm, | |
use_attn_mask=use_sam_mask, | |
save_attn_mask=False, | |
lam_content=1, lam_style=1, | |
sam_mask=sam_mask, use_sam_mask=use_sam_mask, | |
sam_prompt=sam_prompt | |
) | |
for (sampled_c, _, _) in tqdm(sampling_c, total=extras.sampling_configs['timesteps']): | |
sampled_c = sampled_c | |
# Stage B reverse process. | |
with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16): | |
conditions_b['effnet'] = sampled_c | |
unconditions_b['effnet'] = torch.zeros_like(sampled_c) | |
sampling_b = extras_b.gdf.sample( | |
models_b.generator, conditions_b, stage_b_latent_shape, | |
unconditions_b, device=device, **extras_b.sampling_configs, | |
) | |
for (sampled_b, _, _) in tqdm(sampling_b, total=extras_b.sampling_configs['timesteps']): | |
sampled_b = sampled_b | |
sampled = models_b.stage_a.decode(sampled_b).float() | |
sampled = torch.cat([ | |
torch.nn.functional.interpolate(ref_images.cpu(), size=(height, width)), | |
torch.nn.functional.interpolate(ref_style.cpu(), size=(height, width)), | |
sampled.cpu(), | |
], dim=0) | |
# Remove the batch dimension and keep only the generated image | |
sampled = sampled[2] # This selects the generated image, discarding the reference images | |
# Ensure the tensor is in [C, H, W] format | |
if sampled.dim() == 3 and sampled.shape[0] == 3: | |
output_file = 'output_compo.png' | |
sampled_image = T.ToPILImage()(sampled) # Convert tensor to PIL image | |
sampled_image.save(output_file) # Save the image as a PNG | |
else: | |
raise ValueError(f"Expected tensor of shape [3, H, W] but got {sampled.shape}") | |
return output_file # Return the path to the saved image | |
finally: | |
# Reset the state after inference, regardless of success or failure | |
reset_compo_inference_state() | |
def run(style_description, style_reference_image, subject_prompt, subject_reference, use_subject_ref): | |
result = None | |
if use_subject_ref is True: | |
result = infer_compo(style_description, style_reference_image, subject_prompt, subject_reference) | |
else: | |
result = infer(style_reference_image, style_description, subject_prompt) | |
return result | |
import gradio as gr | |
with gr.Blocks() as demo: | |
with gr.Column(): | |
gr.Markdown("# RB-Modulation") | |
gr.Markdown("## Training-Free Personalization of Diffusion Models using Stochastic Optimal Control") | |
gr.HTML(""" | |
<div style="display:flex;column-gap:4px;"> | |
<a href='https://rb-modulation.github.io'> | |
<img src='https://img.shields.io/badge/Project-Page-Green'> | |
</a> | |
<a href='https://arxiv.org/pdf/2405.17401'> | |
<img src='https://img.shields.io/badge/Paper-Arxiv-red'> | |
</a> | |
</div> | |
""") | |
with gr.Row(): | |
with gr.Column(): | |
style_reference_image = gr.Image( | |
label = "Style Reference Image", | |
type = "filepath" | |
) | |
style_description = gr.Textbox( | |
label ="Style Description" | |
) | |
subject_prompt = gr.Textbox( | |
label = "Subject Prompt" | |
) | |
with gr.Accordion("Advanced Settings", open=False): | |
subject_reference = gr.Image(type="filepath") | |
use_subject_ref = gr.Checkbox(label="Use Subject Image as Reference", value=False) | |
submit_btn = gr.Button("Submit") | |
with gr.Column(): | |
output_image = gr.Image(label="Output Image") | |
''' | |
submit_btn.click( | |
fn = infer, | |
inputs = [style_reference_image, style_description, subject_prompt], | |
outputs = [output_image] | |
) | |
''' | |
submit_btn.click( | |
fn = run, | |
inputs = [style_description, style_reference_image, subject_prompt, subject_reference, use_subject_ref], | |
outputs = [output_image] | |
) | |
demo.launch() |