| import sys |
| from pathlib import Path |
| from typing import List, Optional |
|
|
| import gradio as gr |
| import torch |
| from PIL import Image |
| from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler |
| from huggingface_hub import snapshot_download |
| from transformers import CLIPTokenizer |
|
|
| import constants |
| from checkpoint_handler import CheckpointHandler |
| from models.neti_clip_text_encoder import NeTICLIPTextModel |
| from models.xti_attention_processor import XTIAttenProc |
| from prompt_manager import PromptManager |
| from scripts.inference import run_inference |
|
|
| sys.path.append(".") |
| sys.path.append("..") |
|
|
| DESCRIPTION = ''' |
| # A Neural Space-Time Representation for Text-to-Image Personalization |
| <p style="text-align: center;"> |
| This is a demo for our <a href="https://arxiv.org/abs/2305.15391">paper</a>: ''A Neural Space-Time Representation |
| for Text-to-Image Personalization''. |
| <br> |
| Project page and code is available <a href="https://neuraltextualinversion.github.io/NeTI/">here</a>. |
| <br> |
| We introduce a new text-conditioning latent space P* that is dependent on both the denoising process timestep and |
| the U-Net layers. |
| This space-time representation is learned implicitly via a small mapping network. |
| <br> |
| Here, you can generate images using one of the concepts trained in our paper. Simply select your concept and |
| random seed. |
| <br> |
| You can also choose different truncation values to play with the reconstruction vs. editability of the concept. |
| </p> |
| ''' |
|
|
| CONCEPT_TO_PLACEHOLDER = { |
| 'barn': '<barn>', |
| 'cat': '<cat>', |
| 'clock': '<clock>', |
| 'colorful_teapot': '<colorful-teapot>', |
| 'dangling_child': '<dangling-child>', |
| 'dog': '<dog>', |
| 'elephant': '<elephant>', |
| 'fat_stone_bird': '<stone-bird>', |
| 'headless_statue': '<headless-statue>', |
| 'lecun': '<lecun>', |
| 'maeve': '<maeve-dog>', |
| 'metal_bird': '<metal-bird>', |
| 'mugs_skulls': '<mug-skulls>', |
| 'rainbow_cat': '<rainbow-cat>', |
| 'red_bowl': '<red-bowl>', |
| 'teddybear': '<teddybear>', |
| 'tortoise_plushy': '<tortoise-plushy>', |
| 'wooden_pot': '<wooden-pot>' |
| } |
|
|
| MODELS_PATH = Path('./trained_models') |
| MODELS_PATH.mkdir(parents=True, exist_ok=True) |
|
|
|
|
| def load_stable_diffusion_model(pretrained_model_name_or_path: str, |
| num_denoising_steps: int = 50, |
| torch_dtype: torch.dtype = torch.float16) -> StableDiffusionPipeline: |
| tokenizer = CLIPTokenizer.from_pretrained( |
| pretrained_model_name_or_path, subfolder="tokenizer") |
| text_encoder = NeTICLIPTextModel.from_pretrained( |
| pretrained_model_name_or_path, subfolder="text_encoder", torch_dtype=torch_dtype, |
| ) |
| pipeline = StableDiffusionPipeline.from_pretrained( |
| pretrained_model_name_or_path, |
| torch_dtype=torch_dtype, |
| text_encoder=text_encoder, |
| tokenizer=tokenizer |
| ).to("cuda") |
| pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config) |
| pipeline.scheduler.set_timesteps(num_denoising_steps, device=pipeline.device) |
| pipeline.unet.set_attn_processor(XTIAttenProc()) |
| return pipeline |
|
|
|
|
| def get_possible_concepts() -> List[str]: |
| objects = [x for x in MODELS_PATH.iterdir() if x.is_dir()] |
| return [x.name for x in objects] |
|
|
|
|
| def load_sd_and_all_tokens(): |
| mappers = {} |
| pipeline = load_stable_diffusion_model(pretrained_model_name_or_path="CompVis/stable-diffusion-v1-4") |
| print("Downloading all models from HF Hub...") |
| snapshot_download(repo_id="neural-ti/NeTI", local_dir='./trained_models') |
| print("Done.") |
| concepts = get_possible_concepts() |
| for concept in concepts: |
| print(f"Loading model for concept: {concept}") |
| learned_embeds_path = MODELS_PATH / concept / f"{concept}-learned_embeds.bin" |
| mapper_path = MODELS_PATH / concept / f"{concept}-mapper.pt" |
| train_cfg, mapper = CheckpointHandler.load_mapper(mapper_path=mapper_path) |
| placeholder_token, placeholder_token_id = CheckpointHandler.load_learned_embed_in_clip( |
| learned_embeds_path=learned_embeds_path, |
| text_encoder=pipeline.text_encoder, |
| tokenizer=pipeline.tokenizer |
| ) |
| mappers[concept] = { |
| "mapper": mapper, |
| "placeholder_token": placeholder_token, |
| "placeholder_token_id": placeholder_token_id |
| } |
| return mappers, pipeline |
|
|
|
|
| mappers, pipeline = load_sd_and_all_tokens() |
|
|
|
|
| def main_pipeline(concept_name: str, |
| prompt_input: str, |
| seed: int, |
| use_truncation: bool = False, |
| truncation_idx: Optional[int] = None) -> Image.Image: |
| pipeline.text_encoder.text_model.embeddings.set_mapper(mappers[concept_name]["mapper"]) |
| placeholder_token = mappers[concept_name]["placeholder_token"] |
| placeholder_token_id = mappers[concept_name]["placeholder_token_id"] |
| prompt_manager = PromptManager(tokenizer=pipeline.tokenizer, |
| text_encoder=pipeline.text_encoder, |
| timesteps=pipeline.scheduler.timesteps, |
| unet_layers=constants.UNET_LAYERS, |
| placeholder_token=placeholder_token, |
| placeholder_token_id=placeholder_token_id, |
| torch_dtype=torch.float16) |
| image = run_inference(prompt=prompt_input.replace("*", CONCEPT_TO_PLACEHOLDER[concept_name]), |
| pipeline=pipeline, |
| prompt_manager=prompt_manager, |
| seeds=[int(seed)], |
| num_images_per_prompt=1, |
| truncation_idx=truncation_idx if use_truncation else None) |
| return [image] |
|
|
|
|
| with gr.Blocks(css='style.css') as demo: |
| gr.Markdown(DESCRIPTION) |
|
|
| gr.HTML('''<a href="https://huggingface.co/spaces/neural-ti/NeTI?duplicate=true"><img src="https://bit.ly/3gLdBN6" |
| alt="Duplicate Space"></a>''') |
|
|
| with gr.Row(): |
| with gr.Column(): |
| concept = gr.Dropdown(get_possible_concepts(), multiselect=False, label="Concept", |
| info="Choose your concept") |
| prompt = gr.Textbox(label="Input prompt", info="Input prompt with placeholder for concept. " |
| "Please use * to specify the concept.") |
| random_seed = gr.Number(value=42, label="Random seed", precision=0) |
| use_truncation = gr.Checkbox(label="Use inference-time dropout", |
| info="Whether to use our dropout technique when computing the concept " |
| "embeddings.") |
| truncation_idx = gr.Slider(8, 128, label="Truncation index", |
| info="If using truncation, which index to truncate from. Lower numbers tend to " |
| "result in more editable images, but at the cost of reconstruction.") |
| run_button = gr.Button('Generate') |
|
|
| with gr.Column(): |
| result = gr.Gallery(label='Result') |
| inputs = [concept, prompt, random_seed, use_truncation, truncation_idx] |
| outputs = [result] |
| run_button.click(fn=main_pipeline, inputs=inputs, outputs=outputs) |
|
|
| with gr.Row(): |
| examples = [ |
| ["maeve", "A photo of * swimming in the ocean", 5196, True, 16], |
| ["dangling_child", "A photo of * in Times Square", 3552126062741487430, False, 8], |
| ["teddybear", "A photo of * at his graduation ceremony after finishing his PhD", 263, True, 32], |
| ["red_bowl", "A * vase filled with flowers", 13491504810502930872, False, 8], |
| ["metal_bird", "* in a comic book", 1028, True, 24], |
| ["fat_stone_bird", "A movie poster of The Rock, featuring * about on Godzilla", 7393181316156044422, True, |
| 64], |
| ] |
| gr.Examples(examples=examples, |
| inputs=[concept, prompt, random_seed, use_truncation, truncation_idx], |
| outputs=[result], |
| fn=main_pipeline, |
| cache_examples=True) |
|
|
| demo.queue(max_size=50).launch(share=False) |
|
|
