| | import sys |
| | from dataclasses import dataclass |
| | from pathlib import Path |
| | from typing import List, Optional, Union |
| |
|
| | import numpy as np |
| | import torch |
| | from diffusers import DDIMScheduler, DiffusionPipeline |
| | from diffusers.utils import BaseOutput |
| | from PIL import Image |
| |
|
| | _ROOT = Path(__file__).resolve().parent |
| | if str(_ROOT) not in sys.path: |
| | sys.path.insert(0, str(_ROOT)) |
| |
|
| | |
| | sys.modules["pipeline"] = sys.modules[__name__] |
| |
|
| | from modular_pipeline import load_components, resolve_model_root |
| |
|
| |
|
| | @dataclass |
| | class CRSDiffPipelineOutput(BaseOutput): |
| | images: List[Image.Image] |
| |
|
| |
|
| | class CRSDiffPipeline(DiffusionPipeline): |
| | def register_modules(self, **kwargs): |
| | for name, module in kwargs.items(): |
| | if module is None or ( |
| | isinstance(module, (tuple, list)) and len(module) > 0 and module[0] is None |
| | ): |
| | self.register_to_config(**{name: (None, None)}) |
| | setattr(self, name, module) |
| | elif _is_component_list(module): |
| | self.register_to_config(**{name: (module[0], module[1])}) |
| | setattr(self, name, module) |
| | else: |
| | from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple |
| |
|
| | library, class_name = _fetch_class_library_tuple(module) |
| | self.register_to_config(**{name: (library, class_name)}) |
| | setattr(self, name, module) |
| |
|
| | def __init__( |
| | self, |
| | crs_model=None, |
| | scheduler=None, |
| | scale_factor: float = 0.18215, |
| | model_path: Optional[Union[str, Path]] = None, |
| | _name_or_path: Optional[Union[str, Path]] = None, |
| | ): |
| | super().__init__() |
| | if _is_component_list(crs_model) or _is_component_list(scheduler): |
| | model_root = ( |
| | resolve_model_root(model_path) |
| | or resolve_model_root(_name_or_path) |
| | or resolve_model_root(getattr(getattr(self, "config", None), "_name_or_path", None)) |
| | ) |
| | if model_root is None: |
| | raise ValueError( |
| | "CRS-Diff received config placeholders but could not resolve model path. " |
| | "Pass `model_path` or load via DiffusionPipeline.from_pretrained(<path>, custom_pipeline=...)." |
| | ) |
| | loaded = load_components(model_root) |
| | crs_model = loaded["crs_model"] |
| | scheduler = loaded["scheduler"] |
| | scale_factor = loaded["scale_factor"] |
| |
|
| | self.register_modules(crs_model=crs_model, scheduler=scheduler) |
| | self.vae_scale_factor = scale_factor |
| |
|
| | @property |
| | def device(self) -> torch.device: |
| | params = list(self.crs_model.parameters()) |
| | if params: |
| | return params[0].device |
| | return torch.device("cpu") |
| |
|
| | @classmethod |
| | def from_pretrained( |
| | cls, |
| | pretrained_model_name_or_path: Union[str, Path], |
| | device: Optional[Union[str, torch.device]] = None, |
| | subfolder: Optional[str] = None, |
| | **kwargs, |
| | ) -> "CRSDiffPipeline": |
| | path = resolve_model_root(pretrained_model_name_or_path) |
| | if path is None: |
| | raise ValueError(f"Could not resolve CRS-Diff model root from: {pretrained_model_name_or_path}") |
| |
|
| | subfolder = kwargs.pop("subfolder", subfolder) |
| | if subfolder == "scheduler": |
| | return DDIMScheduler.from_pretrained(path, subfolder="scheduler") |
| |
|
| | loaded = load_components(path) |
| | pipe = cls(crs_model=loaded["crs_model"], scheduler=loaded["scheduler"], scale_factor=loaded["scale_factor"]) |
| | if device is not None: |
| | pipe = pipe.to(device) |
| | return pipe |
| |
|
| | def _to_tensor(self, x, device: torch.device, dtype=torch.float32) -> torch.Tensor: |
| | if isinstance(x, np.ndarray): |
| | x = torch.from_numpy(x) |
| | if not isinstance(x, torch.Tensor): |
| | raise TypeError("Expected torch.Tensor or np.ndarray for conditioning inputs.") |
| | return x.to(device=device, dtype=dtype) |
| |
|
| | @torch.no_grad() |
| | def __call__( |
| | self, |
| | prompt: Union[str, List[str]], |
| | local_control, |
| | global_control, |
| | metadata, |
| | negative_prompt: Union[str, List[str]] = "", |
| | num_inference_steps: int = 50, |
| | guidance_scale: float = 7.5, |
| | eta: float = 0.0, |
| | strength: float = 1.0, |
| | global_strength: float = 1.0, |
| | generator: Optional[torch.Generator] = None, |
| | output_type: str = "pil", |
| | ) -> CRSDiffPipelineOutput: |
| | device = self.device |
| | local_control = self._to_tensor(local_control, device=device) |
| | global_control = self._to_tensor(global_control, device=device) |
| | metadata = self._to_tensor(metadata, device=device) |
| |
|
| | batch_size = local_control.shape[0] |
| | if isinstance(prompt, str): |
| | prompt = [prompt] * batch_size |
| | if isinstance(negative_prompt, str): |
| | negative_prompt = [negative_prompt] * batch_size |
| |
|
| | if metadata.dim() == 1: |
| | metadata = metadata.unsqueeze(0).repeat(batch_size, 1) |
| |
|
| | cond = { |
| | "local_control": [local_control], |
| | "c_crossattn": [self.crs_model.get_learned_conditioning(prompt)], |
| | "global_control": [global_control], |
| | } |
| | un_cond = { |
| | "local_control": [local_control], |
| | "c_crossattn": [self.crs_model.get_learned_conditioning(negative_prompt)], |
| | "global_control": [torch.zeros_like(global_control)], |
| | } |
| |
|
| | if hasattr(self.crs_model, "local_control_scales"): |
| | self.crs_model.local_control_scales = [strength] * 13 |
| |
|
| | _, _, h, w = local_control.shape |
| | latents = torch.randn( |
| | (batch_size, self.crs_model.channels, h // 8, w // 8), |
| | generator=generator, |
| | device=device, |
| | ) |
| | latents = latents * self.scheduler.init_noise_sigma |
| |
|
| | self.scheduler.set_timesteps(num_inference_steps, device=device) |
| | for t in self.scheduler.timesteps: |
| | ts = torch.full((batch_size,), int(t), device=device, dtype=torch.long) |
| | if guidance_scale > 1.0: |
| | noise_text = self.crs_model.apply_model(latents, ts, cond, metadata, global_strength) |
| | noise_uncond = self.crs_model.apply_model(latents, ts, un_cond, metadata, global_strength) |
| | noise_pred = noise_uncond + guidance_scale * (noise_text - noise_uncond) |
| | else: |
| | noise_pred = self.crs_model.apply_model(latents, ts, cond, metadata, global_strength) |
| |
|
| | latents = self.scheduler.step( |
| | model_output=noise_pred, |
| | timestep=t, |
| | sample=latents, |
| | eta=eta, |
| | generator=generator, |
| | return_dict=True, |
| | ).prev_sample |
| |
|
| | images = self.crs_model.decode_first_stage(latents) |
| | images = images.clamp(-1, 1) |
| | images = ((images + 1.0) / 2.0).permute(0, 2, 3, 1).cpu().numpy() |
| | images = (images * 255.0).clip(0, 255).astype(np.uint8) |
| |
|
| | if output_type == "pil": |
| | images = [Image.fromarray(img) for img in images] |
| | elif output_type != "numpy": |
| | raise ValueError("output_type must be 'pil' or 'numpy'") |
| |
|
| | return CRSDiffPipelineOutput(images=images) |
| |
|
| |
|
| | def _is_component_list(v): |
| | return ( |
| | isinstance(v, (list, tuple)) |
| | and len(v) == 2 |
| | and isinstance(v[0], str) |
| | and isinstance(v[1], str) |
| | ) |
| |
|
