huggan/sefa

1# python 3.72"""Demo."""3 4import numpy as np5import torch6import streamlit as st7import SessionState8 9from models import parse_gan_type10from utils import to_tensor11from utils import postprocess12from utils import load_generator13from utils import factorize_weight14 15 16@st.cache(allow_output_mutation=True, show_spinner=False)17def get_model(model_name):18    """Gets model by name."""19    return load_generator(model_name, from_hf_hub=True)20 21 22@st.cache(allow_output_mutation=True, show_spinner=False)23def factorize_model(model, layer_idx):24    """Factorizes semantics from target layers of the given model."""25    return factorize_weight(model, layer_idx)26 27 28def sample(model, gan_type, num=1):29    """Samples latent codes."""30    codes = torch.randn(num, model.z_space_dim)31    if gan_type == 'pggan':32        codes = model.layer0.pixel_norm(codes)33    elif gan_type == 'stylegan':34        codes = model.mapping(codes)['w']35        codes = model.truncation(codes,36                                 trunc_psi=0.7,37                                 trunc_layers=8)38    elif gan_type == 'stylegan2':39        codes = model.mapping(codes)['w']40        codes = model.truncation(codes,41                                 trunc_psi=0.5,42                                 trunc_layers=18)43    codes = codes.detach().cpu().numpy()44    return codes45 46 47@st.cache(allow_output_mutation=True, show_spinner=False)48def synthesize(model, gan_type, code):49    """Synthesizes an image with the give code."""50    if gan_type == 'pggan':51        image = model(to_tensor(code))['image']52    elif gan_type in ['stylegan', 'stylegan2']:53        image = model.synthesis(to_tensor(code))['image']54    image = postprocess(image)[0]55    return image56 57def _update_slider():58    num_semantics = st.session_state["num_semantics"]59    for sem_idx in range(num_semantics):60        st.session_state[f"semantic_slider_{sem_idx}"] = 061 62 63"""Main function (loop for StreamLit)."""64st.title('Closed-Form Factorization of Latent Semantics in GANs')65st.markdown("This space is the ported version of [Closed-Form Factorization of Latent Semantics in GANs](https://github.com/genforce/sefa). It reads all sample models from the Hugging Face Hub")66st.markdown("---")67    68st.sidebar.title('Options')69st.sidebar.button('Reset', on_click=_update_slider, kwargs={})70 71model_name = st.sidebar.selectbox(72    'Model to Interpret',73    ['pggan_celebahq1024', 'stylegan_animeface512', 'stylegan_car512', 'stylegan_cat256'])74 75model = get_model(model_name)76gan_type = parse_gan_type(model)77layer_idx = st.sidebar.selectbox(78    'Layers to Interpret',79    ['all', '0-1', '2-5', '6-13'])80layers, boundaries, eigen_values = factorize_model(model, layer_idx)81 82num_semantics = st.sidebar.number_input(83    'Number of semantics', value=5, min_value=0, max_value=None, step=1, key="num_semantics")84steps = {sem_idx: 0 for sem_idx in range(num_semantics)}85if gan_type == 'pggan':86    max_step = 5.087elif gan_type == 'stylegan':88    max_step = 2.089elif gan_type == 'stylegan2':90    max_step = 15.091for sem_idx in steps:92    eigen_value = eigen_values[sem_idx]93    steps[sem_idx] = st.sidebar.slider(94        f'Semantic {sem_idx:03d} (eigen value: {eigen_value:.3f})',95        value=0.0,96        min_value=-max_step,97        max_value=max_step,98        step=0.04 * max_step,99        key=f"semantic_slider_{sem_idx}")100 101image_placeholder = st.empty()102button_placeholder = st.empty()103button_totally_random = st.empty()104 105try:106    base_codes = np.load(f'latent_codes/{model_name}_latents.npy')107except FileNotFoundError:108    base_codes = sample(model, gan_type)109 110state = SessionState.get(model_name=model_name,111                            code_idx=0,112                            codes=base_codes[0:1])113if state.model_name != model_name:114    state.model_name = model_name115    state.code_idx = 0116    state.codes = base_codes[0:1]117 118if button_placeholder.button('Next Sample'):119    state.code_idx += 1120    if state.code_idx < base_codes.shape[0]:121        state.codes = base_codes[state.code_idx][np.newaxis]122    else:123        state.codes = sample(model, gan_type)124 125if button_totally_random.button('Totally Random'):126    state.codes = sample(model, gan_type)127 128code = state.codes.copy()129for sem_idx, step in steps.items():130    if gan_type == 'pggan':131        code += boundaries[sem_idx:sem_idx + 1] * step132    elif gan_type in ['stylegan', 'stylegan2']:133        code[:, layers, :] += boundaries[sem_idx:sem_idx + 1] * step134image = synthesize(model, gan_type, code)135image_placeholder.image(image / 255.0)136 137st.markdown("---")138st.markdown("""This space was created by [johko](https://twitter.com/johko990). Main credits go to the original authors Yujun Shen and Bolei Zhou, who created a great code base to work on. 139            This version loads all models from the Hugging Face Hub.""")140