fix bug for cpu running

app.py

@@ -9,29 +9,18 @@ import time
 import legacy
 import torch
 import glob
-
 import cv2
-import signal
+
 from torch_utils import misc
 from renderer import Renderer
 from training.networks import Generator
 from huggingface_hub import hf_hub_download
 
 
-device = torch.device('cuda')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 port   = int(sys.argv[1]) if len(sys.argv) > 1 else 21111
 
 
-
-def handler(signum, frame):
-    res = input("Ctrl-c was pressed. Do you really want to exit? y/n ")
-    if res == 'y':
-        gr.close_all()
-        exit(1)
-
-signal.signal(signal.SIGINT, handler)
-
-
 def set_random_seed(seed):
     torch.manual_seed(seed)
     np.random.seed(seed)
@@ -202,11 +191,12 @@ yaw    = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="yaw")
 pitch  = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="pitch")
 roll   = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="roll (optional, not suggested)")
 fov    = gr.inputs.Slider(minimum=9, maximum=15, default=12, label="fov")
-css = ".output_image {height: 40rem !important; width: 100% !important;}"
+css = ".output-image, .input-image, .image-preview {height: 600px !important} "
 
 gr.Interface(fn=f_synthesis,
              inputs=[model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, "state"],
              title="Interctive Web Demo for StyleNeRF (ICLR 2022)",
+             description="Demo for ICLR 2022 Papaer: A Style-based 3D-Aware Generator for High-resolution Image Synthesis. Currently the demo runs on CPU only."
              outputs=["image", "state"],
              layout='unaligned',
              css=css, theme='dark-huggingface',

training/networks.py

@@ -794,7 +794,7 @@ class SynthesisBlock(torch.nn.Module):
     def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, add_on=None, block_noise=None, disable_rgb=False, **layer_kwargs):
         misc.assert_shape(ws, [None, self.num_conv + self.num_torgb, self.w_dim])
         w_iter = iter(ws.unbind(dim=1))
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        dtype = torch.float16 if (self.use_fp16 and x.device.type == 'cuda') and not force_fp32 else torch.float32
         memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
         if fused_modconv is None:
             with misc.suppress_tracer_warnings(): # this value will be treated as a constant
@@ -937,7 +937,7 @@ class SynthesisBlock3(torch.nn.Module):
 
     def forward(self, x, img, ws,  force_fp32=False, add_on=None, disable_rgb=False, **layer_kwargs):
         w_iter = iter(ws.unbind(dim=1))
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        dtype = torch.float16 if (self.use_fp16 and x.device.type == 'cuda') and not force_fp32 else torch.float32
         memory_format = torch.contiguous_format
         
         # Main layers.
@@ -1141,7 +1141,7 @@ class DiscriminatorBlock(torch.nn.Module):
                 trainable=next(trainable_iter), resample_filter=resample_filter, channels_last=self.channels_last)
 
     def forward(self, x, img, force_fp32=False, downsampler=None):
-        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
+        dtype = torch.float16 if (self.use_fp16 and x.device.type == 'cuda') and not force_fp32 else torch.float32
         memory_format = torch.channels_last if self.channels_last and not force_fp32 else torch.contiguous_format
 
         # Input.