update interface

interface.py

@@ -5,13 +5,18 @@ import gradio as gr
 import numpy as np
 import torch
 import pickle
+import PIL.Image
 import types
 
+from projector import project, imageio, _MODELS
+
 from huggingface_hub import hf_hub_url, cached_download
 
-# with open('../models/gamma500/network-snapshot-010000.pkl', 'rb') as f:
+# with open("../models/gamma500/network-snapshot-010000.pkl", "rb") as f:
+# with open("../models/gamma400/network-snapshot-010600.pkl", "rb") as f:
+# with open("../models/gamma400/network-snapshot-019600.pkl", "rb") as f:
 with open(cached_download(hf_hub_url('ykilcher/apes', 'gamma500/network-snapshot-010000.pkl')), 'rb') as f:
-    G = pickle.load(f)['G_ema']# torch.nn.Module
+    G = pickle.load(f)["G_ema"]  # torch.nn.Module
 
 device = torch.device("cpu")
 if torch.cuda.is_available():
@@ -35,36 +40,87 @@ else:
     G.synthesis.forward = types.MethodType(_new_synthesis_forward, G.synthesis)
 
 
-def generate(num_images, interpolate):
-    if interpolate:
-        z1 = torch.randn([1, G.z_dim])# latent codes
-        z2 = torch.randn([1, G.z_dim])# latent codes
-        zs = torch.cat([z1 + (z2 - z1) * i / (num_images-1) for i in range(num_images)], 0)
-    else:
-        zs = torch.randn([num_images, G.z_dim])# latent codes
+def generate(
+    target_image_upload,
+    # target_image_webcam,
+    num_steps,
+    seed,
+    learning_rate,
+    model_name,
+    normalize_for_clip,
+    loss_type,
+    regularize_noise_weight,
+    initial_noise_factor,
+):
+    seed = round(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    target_image = target_image_upload
+    # if target_image is None:
+        # target_image = target_image_webcam
+    num_steps = round(num_steps)
+    print(type(target_image))
+    print(target_image.dtype)
+    print(target_image.max())
+    print(target_image.min())
+    print(target_image.shape)
+    target_pil = PIL.Image.fromarray(target_image).convert("RGB")
+    w, h = target_pil.size
+    s = min(w, h)
+    target_pil = target_pil.crop(
+        ((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2)
+    )
+    target_pil = target_pil.resize(
+        (G.img_resolution, G.img_resolution), PIL.Image.LANCZOS
+    )
+    target_uint8 = np.array(target_pil, dtype=np.uint8)
+    target_image = torch.from_numpy(target_uint8.transpose([2, 0, 1])).to(device)
+    projected_w_steps = project(
+        G,
+        target=target_image,
+        num_steps=num_steps,
+        device=device,
+        verbose=True,
+        initial_learning_rate=learning_rate,
+        model_name=model_name,
+        normalize_for_clip=normalize_for_clip,
+        loss_type=loss_type,
+        regularize_noise_weight=regularize_noise_weight,
+        initial_noise_factor=initial_noise_factor,
+    )
     with torch.no_grad():
-        zs = zs.to(device)
-        img = G(zs, None, force_fp32=True, noise_mode='const') 
-        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-    return img.cpu().numpy()
-
-def greet(num_images, interpolate):
-    img = generate(round(num_images), interpolate)
-    imgs = list(img)
-    if len(imgs) == 1:
-        return imgs[0]
-    grid_len = int(np.ceil(np.sqrt(len(imgs)))) * 2
-    grid_height = int(np.ceil(len(imgs) / grid_len))
-    grid = np.zeros((grid_height * imgs[0].shape[0], grid_len * imgs[0].shape[1], 3), dtype=np.uint8)
-    for i, img in enumerate(imgs):
-        y = (i // grid_len) * img.shape[0]
-        x = (i % grid_len) * img.shape[1]
-        grid[y:y+img.shape[0], x:x+img.shape[1], :] = img
-    return grid
+        video = imageio.get_writer(f'proj.mp4', mode='I', fps=10, codec='libx264', bitrate='16M')
+        for w in projected_w_steps:
+            synth_image = G.synthesis(w.to(device).unsqueeze(0), noise_mode="const")
+            synth_image = (synth_image + 1) * (255 / 2)
+            synth_image = (
+                synth_image.permute(0, 2, 3, 1)
+                .clamp(0, 255)
+                .to(torch.uint8)[0]
+                .cpu()
+                .numpy()
+            )
+            video.append_data(np.concatenate([target_uint8, synth_image], axis=1))
+        video.close()
+    return synth_image, "proj.mp4"
 
 
-iface = gr.Interface(fn=greet, inputs=[
-    gr.inputs.Slider(default=1, label="Num Images", minimum=1, maximum=9, step=1),
-    gr.inputs.Checkbox(default=False, label="Interpolate")
-    ], outputs="image")
-iface.launch()
+iface = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.inputs.Image(source="upload", optional=True),
+        # gr.inputs.Image(source="webcam", optional=True),
+        gr.inputs.Number(default=250, label="steps"),
+        gr.inputs.Number(default=69420, label="seed"),
+        gr.inputs.Number(default=0.05, label="learning_rate"),
+        gr.inputs.Dropdown(default='RN50', label="model_name", choices=['vgg16', *_MODELS.keys()]),
+        gr.inputs.Checkbox(default=True, label="normalize_for_clip"),
+        gr.inputs.Dropdown(
+            default="l2", label="loss_type", choices=["l2", "l1", "cosine"]
+        ),
+        gr.inputs.Number(default=1e5, label="regularize_noise_weight"),
+        gr.inputs.Number(default=0.05, label="initial_noise_factor"),
+    ],
+    outputs=["image", "video"],
+)
+iface.launch(inbrowser=True)

interface_old.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+
+import gradio as gr
+
+import numpy as np
+import torch
+import pickle
+import types
+
+from huggingface_hub import hf_hub_url, cached_download
+
+# with open('../models/gamma500/network-snapshot-010000.pkl', 'rb') as f:
+with open(cached_download(hf_hub_url('ykilcher/apes', 'gamma500/network-snapshot-010000.pkl')), 'rb') as f:
+    G = pickle.load(f)['G_ema']# torch.nn.Module
+
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    G = G.to(device)
+else:
+    _old_forward = G.forward
+
+    def _new_forward(self, *args, **kwargs):
+        kwargs["force_fp32"] = True
+        return _old_forward(*args, **kwargs)
+
+    G.forward = types.MethodType(_new_forward, G)
+
+    _old_synthesis_forward = G.synthesis.forward
+
+    def _new_synthesis_forward(self, *args, **kwargs):
+        kwargs["force_fp32"] = True
+        return _old_synthesis_forward(*args, **kwargs)
+
+    G.synthesis.forward = types.MethodType(_new_synthesis_forward, G.synthesis)
+
+
+def generate(num_images, interpolate):
+    if interpolate:
+        z1 = torch.randn([1, G.z_dim])# latent codes
+        z2 = torch.randn([1, G.z_dim])# latent codes
+        zs = torch.cat([z1 + (z2 - z1) * i / (num_images-1) for i in range(num_images)], 0)
+    else:
+        zs = torch.randn([num_images, G.z_dim])# latent codes
+    with torch.no_grad():
+        zs = zs.to(device)
+        img = G(zs, None, force_fp32=True, noise_mode='const') 
+        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    return img.cpu().numpy()
+
+def greet(num_images, interpolate):
+    img = generate(round(num_images), interpolate)
+    imgs = list(img)
+    if len(imgs) == 1:
+        return imgs[0]
+    grid_len = int(np.ceil(np.sqrt(len(imgs)))) * 2
+    grid_height = int(np.ceil(len(imgs) / grid_len))
+    grid = np.zeros((grid_height * imgs[0].shape[0], grid_len * imgs[0].shape[1], 3), dtype=np.uint8)
+    for i, img in enumerate(imgs):
+        y = (i // grid_len) * img.shape[0]
+        x = (i % grid_len) * img.shape[1]
+        grid[y:y+img.shape[0], x:x+img.shape[1], :] = img
+    return grid
+
+
+iface = gr.Interface(fn=greet, inputs=[
+    gr.inputs.Slider(default=1, label="Num Images", minimum=1, maximum=9, step=1),
+    gr.inputs.Checkbox(default=False, label="Interpolate")
+    ], outputs="image")
+iface.launch()

interface_projector.py

@@ -1,126 +0,0 @@
-#!/usr/bin/env python3
-
-import gradio as gr
-
-import numpy as np
-import torch
-import pickle
-import PIL.Image
-import types
-
-from projector import project, imageio, _MODELS
-
-from huggingface_hub import hf_hub_url, cached_download
-
-# with open("../models/gamma500/network-snapshot-010000.pkl", "rb") as f:
-# with open("../models/gamma400/network-snapshot-010600.pkl", "rb") as f:
-# with open("../models/gamma400/network-snapshot-019600.pkl", "rb") as f:
-with open(cached_download(hf_hub_url('ykilcher/apes', 'gamma500/network-snapshot-010000.pkl')), 'rb') as f:
-    G = pickle.load(f)["G_ema"]  # torch.nn.Module
-
-device = torch.device("cpu")
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    G = G.to(device)
-else:
-    _old_forward = G.forward
-
-    def _new_forward(self, *args, **kwargs):
-        kwargs["force_fp32"] = True
-        return _old_forward(*args, **kwargs)
-
-    G.forward = types.MethodType(_new_forward, G)
-
-    _old_synthesis_forward = G.synthesis.forward
-
-    def _new_synthesis_forward(self, *args, **kwargs):
-        kwargs["force_fp32"] = True
-        return _old_synthesis_forward(*args, **kwargs)
-
-    G.synthesis.forward = types.MethodType(_new_synthesis_forward, G.synthesis)
-
-
-def generate(
-    target_image_upload,
-    # target_image_webcam,
-    num_steps,
-    seed,
-    learning_rate,
-    model_name,
-    normalize_for_clip,
-    loss_type,
-    regularize_noise_weight,
-    initial_noise_factor,
-):
-    seed = round(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    target_image = target_image_upload
-    # if target_image is None:
-        # target_image = target_image_webcam
-    num_steps = round(num_steps)
-    print(type(target_image))
-    print(target_image.dtype)
-    print(target_image.max())
-    print(target_image.min())
-    print(target_image.shape)
-    target_pil = PIL.Image.fromarray(target_image).convert("RGB")
-    w, h = target_pil.size
-    s = min(w, h)
-    target_pil = target_pil.crop(
-        ((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2)
-    )
-    target_pil = target_pil.resize(
-        (G.img_resolution, G.img_resolution), PIL.Image.LANCZOS
-    )
-    target_uint8 = np.array(target_pil, dtype=np.uint8)
-    target_image = torch.from_numpy(target_uint8.transpose([2, 0, 1])).to(device)
-    projected_w_steps = project(
-        G,
-        target=target_image,
-        num_steps=num_steps,
-        device=device,
-        verbose=True,
-        initial_learning_rate=learning_rate,
-        model_name=model_name,
-        normalize_for_clip=normalize_for_clip,
-        loss_type=loss_type,
-        regularize_noise_weight=regularize_noise_weight,
-        initial_noise_factor=initial_noise_factor,
-    )
-    with torch.no_grad():
-        video = imageio.get_writer(f'proj.mp4', mode='I', fps=10, codec='libx264', bitrate='16M')
-        for w in projected_w_steps:
-            synth_image = G.synthesis(w.to(device).unsqueeze(0), noise_mode="const")
-            synth_image = (synth_image + 1) * (255 / 2)
-            synth_image = (
-                synth_image.permute(0, 2, 3, 1)
-                .clamp(0, 255)
-                .to(torch.uint8)[0]
-                .cpu()
-                .numpy()
-            )
-            video.append_data(np.concatenate([target_uint8, synth_image], axis=1))
-        video.close()
-    return synth_image, "proj.mp4"
-
-
-iface = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.inputs.Image(source="upload", optional=True),
-        # gr.inputs.Image(source="webcam", optional=True),
-        gr.inputs.Number(default=250, label="steps"),
-        gr.inputs.Number(default=69420, label="seed"),
-        gr.inputs.Number(default=0.05, label="learning_rate"),
-        gr.inputs.Dropdown(default='RN50', label="model_name", choices=['vgg16', *_MODELS.keys()]),
-        gr.inputs.Checkbox(default=True, label="normalize_for_clip"),
-        gr.inputs.Dropdown(
-            default="l2", label="loss_type", choices=["l2", "l1", "cosine"]
-        ),
-        gr.inputs.Number(default=1e5, label="regularize_noise_weight"),
-        gr.inputs.Number(default=0.05, label="initial_noise_factor"),
-    ],
-    outputs=["image", "video"],
-)
-iface.launch(inbrowser=True)