Spaces:
Running on Zero
Running on Zero
| # Copyright (c) 2025 Hansheng Chen | |
| import torch | |
| import inspect | |
| from copy import deepcopy | |
| from mmgen.models.builder import MODELS, build_module | |
| from .base_diffusion import BaseDiffusion | |
| from lakonlab.utils import rgetattr | |
| class LatentDiffusionTextImage(BaseDiffusion): | |
| def __init__(self, | |
| *args, | |
| vae=None, | |
| text_encoder=None, | |
| **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.vae = build_module(vae) if vae is not None else None | |
| self.text_encoder = build_module(text_encoder) if text_encoder is not None else None | |
| def _prepare_train_minibatch_diffusion_args(self, data): | |
| if 'prompt_embed_kwargs' in data: | |
| prompt_embed_kwargs = data['prompt_embed_kwargs'] | |
| elif 'prompt_kwargs' in data: | |
| assert self.text_encoder is not None, 'Text encoder must be provided for encoding text to embeddings.' | |
| prompt_embed_kwargs = self.text_encoder(**data['prompt_kwargs']) | |
| else: | |
| raise ValueError('Either `prompt_embed_kwargs` or `prompt_kwargs` should be provided in the input data.') | |
| if 'latents' in data: | |
| latents = data['latents'] | |
| elif 'images' in data: | |
| assert self.vae is not None, 'VAE must be provided for encoding images to latents.' | |
| with torch.no_grad(): | |
| if hasattr(self.vae, 'dtype'): | |
| vae_dtype = self.vae.dtype | |
| else: | |
| vae_dtype = next(self.vae.parameters()).dtype | |
| latents = self.vae.encode((data['images'] * 2 - 1).to(vae_dtype)).float() | |
| else: | |
| raise ValueError('Either `latents` or `images` should be provided in the input data.') | |
| v = next(iter(prompt_embed_kwargs.values())) | |
| bs = v.size(0) | |
| device = v.device | |
| diffusion_args = (self.patchify(latents), ) | |
| diffusion_kwargs = prompt_embed_kwargs.copy() | |
| distilled_guidance_scale = self.train_cfg.get('distilled_guidance_scale', None) | |
| if distilled_guidance_scale is not None: | |
| distilled_guidance_scale = torch.full( | |
| (bs,), distilled_guidance_scale, dtype=torch.float32, device=device) | |
| diffusion_kwargs.update(guidance=distilled_guidance_scale) | |
| return diffusion_args, diffusion_kwargs, prompt_embed_kwargs, bs, device | |
| def _prepare_train_minibatch_teacher_args(self, data, prompt_embed_kwargs, bs, device): | |
| teacher_guidance_scale = self.train_cfg.get('teacher_guidance_scale', None) | |
| teacher_use_guidance = (teacher_guidance_scale is not None | |
| and teacher_guidance_scale != 0.0 and teacher_guidance_scale != 1.0) | |
| if teacher_use_guidance: | |
| if 'negative_prompt_embed_kwargs' in data: | |
| negative_prompt_embed_kwargs = data['negative_prompt_embed_kwargs'] | |
| elif 'negative_prompt_kwargs' in data: | |
| negative_prompt_embed_kwargs = self.text_encoder(**data['negative_prompt_kwargs']) | |
| else: | |
| raise ValueError( | |
| 'Either `negative_prompt_embed_kwargs` or `negative_prompt_kwargs` should be provided in the ' | |
| 'input data for classifier-free guidance.') | |
| teacher_kwargs = { | |
| k: torch.cat([negative_prompt_embed_kwargs[k], v], dim=0) | |
| for k, v in prompt_embed_kwargs.items()} | |
| teacher_kwargs.update(guidance_scale=teacher_guidance_scale) | |
| else: | |
| teacher_kwargs = prompt_embed_kwargs.copy() | |
| teacher_distilled_guidance_scale = self.train_cfg.get('teacher_distilled_guidance_scale', None) | |
| if teacher_distilled_guidance_scale is not None: | |
| teacher_distilled_guidance_scale = torch.full( | |
| (bs * 2,) if teacher_use_guidance else (bs,), | |
| teacher_distilled_guidance_scale, dtype=torch.float32, device=device) | |
| teacher_kwargs.update(guidance=teacher_distilled_guidance_scale) | |
| return teacher_kwargs | |
| def _prepare_train_minibatch_args(self, data, running_status=None): | |
| diffusion_args, diffusion_kwargs, prompt_embed_kwargs, bs, device = \ | |
| self._prepare_train_minibatch_diffusion_args(data) | |
| parameters = inspect.signature(rgetattr(self.diffusion, 'forward_train')).parameters | |
| if 'running_status' in parameters: | |
| diffusion_kwargs['running_status'] = running_status | |
| if 'teacher' in parameters and 'teacher_kwargs' in parameters and self.teacher is not None: | |
| teacher_kwargs = self._prepare_train_minibatch_teacher_args( | |
| data, prompt_embed_kwargs, bs, device) | |
| diffusion_kwargs.update( | |
| teacher=self.teacher, | |
| teacher_kwargs=teacher_kwargs) | |
| return bs, diffusion_args, diffusion_kwargs | |
| def val_step(self, data, test_cfg_override=dict(), **kwargs): | |
| if 'prompt_embed_kwargs' in data: | |
| prompt_embed_kwargs = data['prompt_embed_kwargs'] | |
| elif 'prompt_kwargs' in data: | |
| assert self.text_encoder is not None, 'Text encoder must be provided for encoding text to embeddings.' | |
| prompt_embed_kwargs = self.text_encoder(**data['prompt_kwargs']) | |
| else: | |
| raise ValueError('Either `prompt_embed_kwargs` or `prompt_kwargs` should be provided in the input data.') | |
| v = next(iter(prompt_embed_kwargs.values())) | |
| bs = v.size(0) | |
| device = v.device | |
| cfg = deepcopy(self.test_cfg) | |
| cfg.update(test_cfg_override) | |
| guidance_scale = cfg.get('guidance_scale', 1.0) | |
| diffusion = self.diffusion_ema if self.diffusion_use_ema else self.diffusion | |
| with torch.no_grad(): | |
| use_guidance = guidance_scale != 0.0 and guidance_scale != 1.0 | |
| if use_guidance: | |
| if 'negative_prompt_embed_kwargs' in data: | |
| negative_prompt_embed_kwargs = data['negative_prompt_embed_kwargs'] | |
| elif 'negative_prompt_kwargs' in data: | |
| negative_prompt_embed_kwargs = self.text_encoder(**data['negative_prompt_kwargs']) | |
| else: | |
| raise ValueError( | |
| 'Either `negative_prompt_embed_kwargs` or `negative_prompt_kwargs` should be provided in the ' | |
| 'input data for classifier-free guidance.') | |
| kwargs = { | |
| k: torch.cat([negative_prompt_embed_kwargs[k], v], dim=0) | |
| for k, v in prompt_embed_kwargs.items()} | |
| else: | |
| kwargs = prompt_embed_kwargs.copy() | |
| distilled_guidance_scale = cfg.get('distilled_guidance_scale', None) | |
| if distilled_guidance_scale is not None: | |
| distilled_guidance_scale = torch.full( | |
| (bs * 2,) if use_guidance else (bs,), | |
| distilled_guidance_scale, dtype=torch.float32, device=device) | |
| kwargs.update(guidance=distilled_guidance_scale) | |
| if 'noise' in data: | |
| noise = data['noise'] | |
| else: | |
| latent_size = cfg['latent_size'] | |
| noise = torch.randn((bs, *latent_size), device=device) | |
| noise = self.patchify(noise) | |
| latents_out = diffusion( | |
| noise=noise, | |
| guidance_scale=guidance_scale, | |
| test_cfg_override=test_cfg_override, | |
| **kwargs) | |
| latents_out = self.unpatchify(latents_out) | |
| if hasattr(self.vae, 'dtype'): | |
| vae_dtype = self.vae.dtype | |
| else: | |
| vae_dtype = next(self.vae.parameters()).dtype | |
| latents_out = latents_out.to(vae_dtype) | |
| out_images = (self.vae.decode(latents_out).float() / 2 + 0.5).clamp(min=0, max=1) | |
| return dict(num_samples=bs, pred_imgs=out_images) | |