Hugging Face's logo

Spaces:
skimai
/
DragGAN_Streamlit
Runtime error

App Files Community
DragGAN_Streamlit
File size: 6,948 Bytes 
da180b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
 
import time
import torch
import streamlit as st
from PIL import Image, ImageDraw
from streamlit_image_coordinates import streamlit_image_coordinates

import draggan
import utils


## Default to CPU if no GPU is available
if torch.cuda.is_available():
  device = torch.device("cuda")
else:
    device = torch.device("cpu")


### Streamlit setup ###

st.set_page_config(
    page_title="DragGAN Demo",
    page_icon="🐉",
    layout="wide",
)


message_container = st.empty()

col1, col2 = st.columns([1, 2])

def reset():
    st.session_state.clear()

def reset_rerun():
    reset()
    st.experimental_rerun()


### Run/Reset buttons in right col ###
with col2:
    st.markdown("")
    but_col1, but_col2 = st.columns([1,7])
    run_button = but_col1.button("▶️ Run")
    reset_button = but_col2.button("🔁 Reset")


### Settings panel in left col ###
with col1:
    st.header("🐉 DragGAN")

    settings_col1, settings_col2 = st.columns([1,1])
    # Models from Self-Distilled SG https://github.com/self-distilled-stylegan/self-distilled-internet-photos
    model_options = {
        "Lions": "https://storage.googleapis.com/self-distilled-stylegan/lions_512_pytorch.pkl",
        "Faces (FFHQ)": "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl",
        "Elephants": "https://storage.googleapis.com/self-distilled-stylegan/elephants_512_pytorch.pkl",
        "Parrots": "https://storage.googleapis.com/self-distilled-stylegan/parrots_512_pytorch.pkl",
        "Horses": "https://storage.googleapis.com/self-distilled-stylegan/horses_256_pytorch.pkl",
        "Bicycles": "https://storage.googleapis.com/self-distilled-stylegan/bicycles_256_pytorch.pkl",
        "Giraffes": "https://storage.googleapis.com/self-distilled-stylegan/giraffes_512_pytorch.pkl",
        "Dogs (1)": "https://storage.googleapis.com/self-distilled-stylegan/dogs_1024_pytorch.pkl",
        "Dogs (2)": "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/afhqdog.pkl",
        "Cats": "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/afhqcat.pkl",
        "Wildlife": "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/afhqwild.pkl",
        "MetFaces": "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl",
    }
    model_name = str(settings_col1.selectbox("Model", list(model_options.keys()), on_change=reset, help="StyleGAN2 model to use, downloaded and cached on first run"))
    model_url = model_options[model_name]
    seed = settings_col2.number_input("Seed", value=22, step=1, min_value=0, on_change=reset, help="Random seed for generating W+ latent")
    target_resolution = int(settings_col1.selectbox("Resolution", [256, 512, 1024], index=1, on_change=reset, help="Resize generated image to this resolution (may be different than native model resolution)"))
    n_iter = int(settings_col1.number_input("Iterations", value=200, step=5, help="Number of iterations to run optimization", on_change=reset))
    step_size = settings_col2.number_input("Step Size", value=1e-3, step=1e-4, min_value=1e-4, format="%.4f", help="Step size (Learning Rate) for gradient descent")
    multiplier = settings_col1.number_input("Speed", value=1.0, step=0.05, min_value=0.05, help="Multiplier for target patch movement")
    tolerance = settings_col2.number_input("Tolerance", value=2, step=1, min_value=1, help="Number of pixels away from target to stop")
    
    display_every = settings_col2.number_input("Display Every", value=25, step=1, min_value=1, help="Display image during optimization every n iterations")
    truncation_psi = settings_col1.number_input("Truncation", value=0.8, step=0.1, min_value=0.0, on_change=reset, help="Truncation trick value to control diversity (higher = more diverse)")
    truncation_cutoff = settings_col2.number_input(
        "Truncation Cutoff", value=8, step=1, min_value=-1, max_value=18, on_change=reset, help="Number of layers to apply truncation to (-1 = all layers)"
    )    
    

    if reset_button:
        reset_rerun()

if "points" not in st.session_state:
    st.session_state["points"] = []
    st.session_state["points_types"] = []
    # State variable to track whether the next click should be a 'handle' or 'target'
    st.session_state["next_click"] = "handle"


s = time.perf_counter()
G = draggan.load_model(model_url, device=device)

if "W" not in st.session_state:
    W = draggan.generate_W(
        G,
        seed=int(seed),
        truncation_psi=truncation_psi,
        truncation_cutoff=int(truncation_cutoff),
        network_pkl=model_url,
        device=device,
    )
else:
    W = st.session_state["W"]

img, F0 = draggan.generate_image(W, G, network_pkl=model_url, device=device)
if img.size[0] != target_resolution:
    img = img.resize((target_resolution, target_resolution))
print(f"Generated image in {(time.perf_counter() - s)*1000:.0f}ms")

# Draw an ellipse at each coordinate in points
if "points" in st.session_state and "points_types" in st.session_state:
    handles, targets = [], []
    for point, point_type in zip(
        st.session_state["points"], st.session_state["points_types"]
    ):
        if point_type == "handle":
            handles.append(point)
        else:
            targets.append(point)
    if len(handles) > 0:
        utils.draw_handle_target_points(img, handles, targets)


### Right column image container ###
with col2:
    empty = st.empty()
    with empty.container():
        value = streamlit_image_coordinates(img, key="pil")
        # New point is clicked
        if value is not None:
            point = value["x"], value["y"]
            if point not in st.session_state["points"]:
                # st.session_state["points"].append(point)
                st.session_state["points"].append(point)
                st.session_state["points_types"].append(st.session_state["next_click"])
                st.session_state["next_click"] = (
                    "target" if st.session_state["next_click"] == "handle" else "handle"
                )
                
                st.experimental_rerun()

## Optimization loop
if run_button:
    if len(handles) > 0 and len(targets) > 0 and len(handles) == len(targets) and all(targets):
        W = draggan.optimize(
            W,
            G,
            handle_points=handles,
            target_points=targets,
            r1=3,
            r2=12,
            tolerance=tolerance,
            max_iter=n_iter,
            lr=step_size,
            multiplier=multiplier,
            empty=empty,
            display_every=display_every,
            target_resolution=target_resolution,
            device=device,
        )
        # st.write(handles)
        # st.write(targets)

        st.session_state.clear()
        st.session_state["W"] = W
        st.experimental_rerun()
    else:
        message_container.warning("Please add at least one handle and one target point.")