Spaces:
Running on Zero
Running on Zero
| """One resident pipeline; original backbone tensors are never cumulatively merged.""" | |
| from pathlib import Path | |
| import torch | |
| from anima_role_ref_kv.components import load_transformer, PromptEncoder, VAE, NEGATIVE | |
| from anima_role_ref_kv.pipeline import sample, resize_reference | |
| from anima_role_ref_kv.weights import load, read_state, swap | |
| from checkpoints import native_plan, diffusers_plan, copy_plan | |
| class Runtime: | |
| def __init__(self, base_dir, adapter_file, device="cuda", adapter_key="v2.0"): | |
| self.base_dir = str(base_dir) | |
| self.net = load_transformer(base_dir, device) | |
| self.text = PromptEncoder(base_dir, device) | |
| self.vae = VAE(base_dir, device) | |
| # Capture references before LoRALinear wraps linears. The adapter keeps | |
| # its frozen base outside its registered parameters; net.state_dict() | |
| # after injection would omit those base weights. | |
| self.targets = {("transformer", k): v for k, v in self.net.state_dict(keep_vars=True).items()} | |
| self.targets.update({("conditioner", k): v for k, v in | |
| self.text.components.text_conditioner.state_dict(keep_vars=True).items()}) | |
| self.base_plan = diffusers_plan(base_dir, self.targets) | |
| self.adapter_states = {adapter_key: read_state(self.net, adapter_file)} | |
| self.adapter = load(self.net, adapter_file) | |
| self.current = "base-v1.0" | |
| self.current_adapter = adapter_key | |
| def select(self, model_key, filename=None): | |
| if self.current == model_key: | |
| return | |
| plan = (self.base_plan if model_key == "base-v1.0" else | |
| (diffusers_plan(filename, self.targets) if Path(filename).is_dir() else | |
| native_plan(filename, self.targets))) | |
| self.current = None # A failed copy must force a complete reload next time. | |
| copy_plan(plan, self.targets) | |
| self.current = model_key | |
| def select_adapter(self, adapter_key, filename=None): | |
| if self.current_adapter == adapter_key: | |
| return | |
| state = self.adapter_states.get(adapter_key) | |
| if state is None: | |
| if not filename: | |
| raise ValueError(f"Adapter file required for {adapter_key}") | |
| state = read_state(self.net, filename) | |
| self.adapter_states[adapter_key] = state | |
| self.current_adapter = None # A failed copy must force a complete reload next time. | |
| swap(self.adapter, state) | |
| self.current_adapter = adapter_key | |
| def generate(self, model_key, filename, reference, prompt, *, adapter_key=None, | |
| adapter_file=None, strength=1.0, negative_prompt=NEGATIVE, | |
| progress_callback=None, **settings): | |
| self.select(model_key, filename) | |
| if adapter_key is not None: | |
| self.select_adapter(adapter_key, adapter_file) | |
| for layer in self.adapter.layers.values(): | |
| layer.multiplier = strength | |
| cfg = settings.get("cfg", 4.) | |
| positive = self.text([prompt]) | |
| negative = self.text([negative_prompt]) if cfg != 1 else None | |
| ref = self.vae.encode(resize_reference(reference)) | |
| latent = sample(self.net, positive, negative, ref, | |
| progress_callback=progress_callback, **settings) | |
| if not torch.isfinite(latent).all(): | |
| raise FloatingPointError("The model produced non-finite latents") | |
| return self.vae.decode(latent) | |