Better colab compat

app.py

@@ -1,10 +1,13 @@
 import os
 import gradio as gr
 import random
+import util
 from random import randint
 import sys
 from subprocess import call
 
+is_colab = util.is_google_colab()
+
 def run_cmd(command):
     try:
         print(command)
@@ -29,8 +32,10 @@ def inference(img):
 input_image = gr.inputs.Image(type="pil", label="Input")
 output_image = gr.outputs.Image(type="file", label="Output")
 
-gr.Interface(
+demo = gr.Interface(
     inference,
     inputs=[input_image],
     outputs=[output_image]
-).launch()
+)
+
+demo.launch(debug=is_colab, share=is_colab)

test.py

@@ -6,18 +6,21 @@ import numpy as np
 import torch
 import architecture as arch
 import multiprocessing
+import util
 
-cpu_count = multiprocessing.cpu_count()
-
-print(cpu_count)
+def is_cuda():
+    if torch.cuda.is_available() or not util.is_google_colab():
+        return True
+    else:
+        return False
 
 model_path = '4x_eula_digimanga_bw_v2_nc1_307k.pth'
 img_path = sys.argv[1]
 output_dir = sys.argv[2]
-device = torch.device('cpu')
+device = torch.device('cuda' if is_cuda() else 'cpu')
 
 model = arch.RRDB_Net(1, 1, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=True)
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cuda' if is_cuda() else 'cpu')), strict=True)
 model.eval()
 
 for k, v in model.named_parameters():

util.py

@@ -0,0 +1,6 @@
+def is_google_colab():
+    try:
+        import google.colab
+        return True
+    except:
+        return False