20f5592077
Co-authored-by: zty624 <164896039+zty624@users.noreply.github.com> Co-authored-by: yhzx233 <16619108+yhzx233@users.noreply.github.com> Co-authored-by: Monalissaa <50078799+Monalissaa@users.noreply.github.com>
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""Application device placement, optional compilation, and progress reporting."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import diffusers
|
|
import torch
|
|
from huggingface_hub.utils import validate_repo_id
|
|
|
|
DEFAULT_MODEL = "OpenMOSS-Team/MOSS-Image-Base"
|
|
|
|
|
|
def model_source(model):
|
|
"""Resolve local bundles without turning Hub IDs into filesystem paths."""
|
|
model = str(model).strip()
|
|
path = Path(model).expanduser()
|
|
if path.is_dir():
|
|
if not (path / "pipeline_moss_image.py").is_file():
|
|
raise ValueError("The local bundle must include pipeline_moss_image.py from the HF repository.")
|
|
return str(path.resolve())
|
|
if model.startswith(("/", ".", "~")):
|
|
raise FileNotFoundError(f"Model directory does not exist: {model}")
|
|
validate_repo_id(model)
|
|
return model
|
|
|
|
|
|
def load_pipeline(model, local_files_only=False):
|
|
return diffusers.DiffusionPipeline.from_pretrained(
|
|
model_source(model),
|
|
custom_pipeline="pipeline_moss_image",
|
|
trust_remote_code=True,
|
|
local_files_only=local_files_only,
|
|
torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
|
|
)
|
|
|
|
|
|
def set_scheduler(pipe, shift):
|
|
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift)
|
|
|
|
|
|
def cuda_budget_gib():
|
|
if not torch.cuda.is_available():
|
|
return 0.0
|
|
total = torch.cuda.get_device_properties(0).total_memory
|
|
return total * torch.cuda.get_per_process_memory_fraction(0) / 1024**3
|
|
|
|
|
|
def configure_pipeline(pipe, offload="auto", compile_transformer=False):
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
if offload == "auto":
|
|
offload = "none" if device == "cpu" or cuda_budget_gib() >= 40 else "model"
|
|
if compile_transformer and offload == "sequential":
|
|
raise ValueError("Compilation is incompatible with sequential offload.")
|
|
if offload == "model":
|
|
pipe.enable_model_cpu_offload(device=device)
|
|
elif offload == "sequential":
|
|
pipe.enable_sequential_cpu_offload(device=device)
|
|
elif offload == "none":
|
|
pipe.to(device)
|
|
else:
|
|
raise ValueError(f"Unknown offload mode: {offload}")
|
|
if compile_transformer:
|
|
for block in (*pipe.transformer.layers, *pipe.transformer.text_embedder.blocks):
|
|
block.compile(backend="inductor", fullgraph=True, dynamic=True, mode="default")
|
|
return pipe
|
|
|
|
|
|
def generate_with_progress(pipe, report_stage, report_step, **kwargs):
|
|
progress_bar = pipe.progress_bar
|
|
|
|
def progress(items):
|
|
report_stage("denoise")
|
|
for index, timestep in enumerate(progress_bar(items)):
|
|
yield timestep
|
|
report_step(pipe, index, timestep, {})
|
|
if kwargs.get("output_type", "pil") != "latent":
|
|
report_stage("decode")
|
|
|
|
pipe.progress_bar = progress
|
|
try:
|
|
report_stage("encode")
|
|
return pipe(**kwargs)
|
|
finally:
|
|
pipe.progress_bar = progress_bar
|