Upload app.py

app.py

@@ -1,7 +1,7 @@
 import sys
 import os
 
-os.system("git clone https://github.com/dunbar12138/pix2pix3D.git")
+os.system("https://github.com/dunbar12138/pix2pix3D.git")
 sys.path.append("pix2pix3D")
 
 from typing import List, Optional, Tuple, Union
@@ -23,6 +23,15 @@ from tqdm import tqdm
 import imageio
 import trimesh
 import mcubes
+import copy
+
+import pickle
+import numpy as np
+import torch
+import dnnlib
+from torch_utils import misc
+from legacy import *
+import io
 
 os.environ["PYOPENGL_PLATFORM"] = "egl"
 
@@ -112,12 +121,59 @@ models={"seg2cat":network_cat
 
 device='cuda' if torch.cuda.is_available() else 'cpu'
 outdir="/content/"
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'torch.storage' and name == '_load_from_bytes':
+            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
+        return super().find_class(module, name)
+
+def load_network_pkl_cpu(f, force_fp16=False):
+    data = CPU_Unpickler(f).load()
+
+    # Legacy TensorFlow pickle => convert.
+    if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):
+        tf_G, tf_D, tf_Gs = data
+        G = convert_tf_generator(tf_G)
+        D = convert_tf_discriminator(tf_D)
+        G_ema = convert_tf_generator(tf_Gs)
+        data = dict(G=G, D=D, G_ema=G_ema)
+
+    # Add missing fields.
+    if 'training_set_kwargs' not in data:
+        data['training_set_kwargs'] = None
+    if 'augment_pipe' not in data:
+        data['augment_pipe'] = None
+
+    # Validate contents.
+    assert isinstance(data['G'], torch.nn.Module)
+    assert isinstance(data['D'], torch.nn.Module)
+    assert isinstance(data['G_ema'], torch.nn.Module)
+    assert isinstance(data['training_set_kwargs'], (dict, type(None)))
+    assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))
+
+    # Force FP16.
+    if force_fp16:
+        for key in ['G', 'D', 'G_ema']:
+            old = data[key]
+            kwargs = copy.deepcopy(old.init_kwargs)
+            fp16_kwargs = kwargs.get('synthesis_kwargs', kwargs)
+            fp16_kwargs.num_fp16_res = 4
+            fp16_kwargs.conv_clamp = 256
+            if kwargs != old.init_kwargs:
+                new = type(old)(**kwargs).eval().requires_grad_(False)
+                misc.copy_params_and_buffers(old, new, require_all=True)
+                data[key] = new
+    return data
 
 def get_all(cfg,input,truncation_psi,mesh_resolution,random_seed,fps,num_frames):
 
         network=models[cfg]
 
-        with dnnlib.util.open_url(network) as f:
+        if device=="cpu":
+          with dnnlib.util.open_url(network) as f:
+              G = load_network_pkl_cpu(f)['G_ema'].eval().to(device)
+        else:    
+          with dnnlib.util.open_url(network) as f:
                 G = legacy.load_network_pkl(f)['G_ema'].eval().to(device)
 
         if cfg == 'seg2cat' or cfg == 'seg2face':
@@ -239,5 +295,6 @@ demo_app = gr.Interface(
     theme="huggingface",
     description=desc,
     examples=examples,
+    cache_examples=True,
 )
 demo_app.launch(debug=True, enable_queue=True)