Add the implementation of SeFa and the interface.

Author: ShenYujun

README.md

@@ -30,6 +30,38 @@ In this repository, we propose a *closed-form* approach, termed as **SeFa**, for
 | Orientation | Vertical Position | Shape
 | ![image](./docs/assets/stylegan_car_orientation.gif) | ![image](./docs/assets/stylegan_car_vertical_position.gif) | ![image](./docs/assets/stylegan_car_shape.gif)
 
+## Semantic Discovery
+
+It is very simple to interpret a particular model with
+
+```bash
+MODEL_NAME=stylegan_animeface512
+LAYER_IDX=0-1
+NUM_SAMPLES=5
+NUM_SEMANTICS=5
+python sefa.py ${MODEL_NAME} \
+    -L ${LAYER_IDX} \
+    -N ${NUM_SAMPLES} \
+    -K ${NUM_SEMANTICS}
+```
+
+After the program finishes, there will be two visualization pages in the directory `results`.
+
+**NOTE:** The pre-trained models are borrowed from the [genforce](https://github.com/genforce/genforce) repository.
+
+## Interface
+
+We also provide an interface for interactive editing based on [StreamLit](https://www.streamlit.io/). This interface can be locally launched with
+
+```bash
+pip install streamlit
+CUDA_VISIBLE_DEVICES=0 streamlit run interface.py
+```
+
+After the interface is launched, users can play with it via a browser.
+
+**NOTE:** We have prepared some latent codes in the directory `latent_codes` to ensure the synthesis quality, which is completely determined by the pre-trained generator. Users can simply skip these prepared codes by clicking the `Random` button.
+
 ## BibTeX
 
 ```bibtex
@@ -40,5 +72,3 @@ In this repository, we propose a *closed-form* approach, termed as **SeFa**, for
   year    = {2020}
 }
 ```
-
-## Code Coming Soon

SessionState.py

@@ -0,0 +1,129 @@
+"""Adds pre-session state to StreamLit.
+
+This file is borrowed from
+https://gist.github.com/tvst/036da038ab3e999a64497f42de966a92
+"""
+
+# pylint: disable=protected-access
+
+try:
+    import streamlit.ReportThread as ReportThread
+    from streamlit.server.Server import Server
+except ModuleNotFoundError:
+    # Streamlit >= 0.65.0
+    import streamlit.report_thread as ReportThread
+    from streamlit.server.server import Server
+
+
+class SessionState(object):
+    """Hack to add per-session state to Streamlit.
+
+    Usage
+    -----
+
+    >>> import SessionState
+    >>>
+    >>> session_state = SessionState.get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    ''
+    >>> session_state.user_name = 'Mary'
+    >>> session_state.favorite_color
+    'black'
+
+    Since you set user_name above, next time your script runs this will be the
+    result:
+    >>> session_state = get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    'Mary'
+
+    """
+
+    def __init__(self, **kwargs):
+        """A new SessionState object.
+
+        Parameters
+        ----------
+        **kwargs : any
+            Default values for the session state.
+
+        Example
+        -------
+        >>> session_state = SessionState(user_name='', favorite_color='black')
+        >>> session_state.user_name = 'Mary'
+        ''
+        >>> session_state.favorite_color
+        'black'
+
+        """
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+
+def get(**kwargs):
+    """Gets a SessionState object for the current session.
+
+    Creates a new object if necessary.
+
+    Parameters
+    ----------
+    **kwargs : any
+        Default values you want to add to the session state, if we're creating a
+        new one.
+
+    Example
+    -------
+    >>> session_state = get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    ''
+    >>> session_state.user_name = 'Mary'
+    >>> session_state.favorite_color
+    'black'
+
+    Since you set user_name above, next time your script runs this will be the
+    result:
+    >>> session_state = get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    'Mary'
+
+    """
+    # Hack to get the session object from Streamlit.
+
+    ctx = ReportThread.get_report_ctx()
+
+    this_session = None
+
+    current_server = Server.get_current()
+    if hasattr(current_server, '_session_infos'):
+        # Streamlit < 0.56
+        session_infos = Server.get_current()._session_infos.values()
+    else:
+        session_infos = Server.get_current()._session_info_by_id.values()
+
+    for session_info in session_infos:
+        s = session_info.session
+        if (
+            # Streamlit < 0.54.0
+            (hasattr(s, '_main_dg') and s._main_dg == ctx.main_dg)
+            or
+            # Streamlit >= 0.54.0
+            (not hasattr(s, '_main_dg') and s.enqueue == ctx.enqueue)
+            or
+            # Streamlit >= 0.65.2
+            (not hasattr(s, '_main_dg') and
+             s._uploaded_file_mgr == ctx.uploaded_file_mgr)
+        ):
+            this_session = s
+
+    if this_session is None:
+        raise RuntimeError(
+            "Oh noes. Couldn't get your Streamlit Session object. "
+            'Are you doing something fancy with threads?')
+
+    # Got the session object! Now let's attach some state into it.
+
+    if not hasattr(this_session, '_custom_session_state'):
+        this_session._custom_session_state = SessionState(**kwargs)
+
+    return this_session._custom_session_state
+
+# pylint: enable=protected-access

interface.py

@@ -0,0 +1,128 @@
+# python 3.7
+"""Demo."""
+
+import numpy as np
+import torch
+import streamlit as st
+import SessionState
+
+from models import parse_gan_type
+from utils import to_tensor
+from utils import postprocess
+from utils import load_generator
+from utils import factorize_weight
+
+
+@st.cache(allow_output_mutation=True, show_spinner=False)
+def get_model(model_name):
+    """Gets model by name."""
+    return load_generator(model_name)
+
+
+@st.cache(allow_output_mutation=True, show_spinner=False)
+def factorize_model(model, layer_idx):
+    """Factorizes semantics from target layers of the given model."""
+    return factorize_weight(model, layer_idx)
+
+
+def sample(model, gan_type, num=1):
+    """Samples latent codes."""
+    codes = torch.randn(num, model.z_space_dim).cuda()
+    if gan_type == 'pggan':
+        codes = model.layer0.pixel_norm(codes)
+    elif gan_type == 'stylegan':
+        codes = model.mapping(codes)['w']
+        codes = model.truncation(codes,
+                                 trunc_psi=0.7,
+                                 trunc_layers=8)
+    elif gan_type == 'stylegan2':
+        codes = model.mapping(codes)['w']
+        codes = model.truncation(codes,
+                                 trunc_psi=0.5,
+                                 trunc_layers=18)
+    codes = codes.detach().cpu().numpy()
+    return codes
+
+
+@st.cache(allow_output_mutation=True, show_spinner=False)
+def synthesize(model, gan_type, code):
+    """Synthesizes an image with the give code."""
+    if gan_type == 'pggan':
+        image = model(to_tensor(code))['image']
+    elif gan_type in ['stylegan', 'stylegan2']:
+        image = model.synthesis(to_tensor(code))['image']
+    image = postprocess(image)[0]
+    return image
+
+
+def main():
+    """Main function (loop for StreamLit)."""
+    st.title('Closed-Form Factorization of Latent Semantics in GANs')
+    st.sidebar.title('Options')
+    reset = st.sidebar.button('Reset')
+
+    model_name = st.sidebar.selectbox(
+        'Model to Interpret',
+        ['stylegan_animeface512', 'stylegan_car512', 'stylegan_cat512',
+         'pggan_celebahq1024'])
+
+    model = get_model(model_name)
+    gan_type = parse_gan_type(model)
+    layer_idx = st.sidebar.selectbox(
+        'Layers to Interpret',
+        ['all', '0-1', '2-5', '6-13'])
+    layers, boundaries, eigen_values = factorize_model(model, layer_idx)
+
+    num_semantics = st.sidebar.number_input(
+        'Number of semantics', value=10, min_value=0, max_value=None, step=1)
+    steps = {sem_idx: 0 for sem_idx in range(num_semantics)}
+    if gan_type == 'pggan':
+        max_step = 5.0
+    elif gan_type == 'stylegan':
+        max_step = 2.0
+    elif gan_type == 'stylegan2':
+        max_step = 15.0
+    for sem_idx in steps:
+        eigen_value = eigen_values[sem_idx]
+        steps[sem_idx] = st.sidebar.slider(
+            f'Semantic {sem_idx:03d} (eigen value: {eigen_value:.3f})',
+            value=0.0,
+            min_value=-max_step,
+            max_value=max_step,
+            step=0.04 * max_step if not reset else 0.0)
+
+    image_placeholder = st.empty()
+    button_placeholder = st.empty()
+
+    try:
+        base_codes = np.load(f'latent_codes/{model_name}_latents.npy')
+    except FileNotFoundError:
+        base_codes = sample(model, gan_type)
+
+    state = SessionState.get(model_name=model_name,
+                             code_idx=0,
+                             codes=base_codes[0:1])
+    if state.model_name != model_name:
+        state.model_name = model_name
+        state.code_idx = 0
+        state.codes = base_codes[0:1]
+
+    if button_placeholder.button('Random', key=0):
+        state.code_idx += 1
+        if state.code_idx < base_codes.shape[0]:
+            state.codes = base_codes[state.code_idx][np.newaxis]
+        else:
+            state.codes = sample(model, gan_type)
+
+    code = state.codes.copy()
+    for sem_idx, step in steps.items():
+        if gan_type == 'pggan':
+            code += boundaries[sem_idx:sem_idx + 1] * step
+        elif gan_type in ['stylegan', 'stylegan2']:
+            code[:, layers, :] += boundaries[sem_idx:sem_idx + 1] * step
+    image = synthesize(model, gan_type, code)
+    image_placeholder.image(image / 255.0)
+
+
+if __name__ == '__main__':
+    main()