add new style

.gitattributes

@@ -26,3 +26,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 model.pth filter=lfs diff=lfs merge=lfs -text
+model2.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -83,18 +83,26 @@ class Generator(nn.Module):
 
         return out
 
-model = Generator(3, 1, 3)
-model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
-model.eval()
+model1 = Generator(3, 1, 3)
+model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
+model1.eval()
 
-def predict(input_img):
+model2 = Generator(3, 1, 3)
+model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
+model2.eval()
+
+def predict(input_img, ver):
     input_img = Image.open(input_img)
     transform = transforms.Compose([transforms.Resize(256, Image.BICUBIC), transforms.ToTensor()])
     input_img = transform(input_img)
     input_img = torch.unsqueeze(input_img, 0)
 
+    drawing = 0
     with torch.no_grad():
-        drawing = model(input_img)[0].detach()
+        if ver == 'style 2':
+            drawing = model2(input_img)[0].detach()
+        else:
+            drawing = model1(input_img)[0].detach()
     
     drawing = transforms.ToPILImage()(drawing)
     return drawing
@@ -102,9 +110,11 @@ def predict(input_img):
 title="informative-drawings"
 description="Gradio Demo for line drawing generation. "
 # article = "<p style='text-align: center'><a href='TODO' target='_blank'>Project Page</a> | <a href='codelink' target='_blank'>Github</a></p>"
-examples=[['cat.png'], ['lizard.png'], ['bridge.png']]
+examples=[['cat.png', 'style 1'], ['bridge.png', 'style 1'], ['lizard.png', 'style 2'],]
 
 
-iface = gr.Interface(predict, gr.inputs.Image(type='filepath'), "image", title=title,description=description,examples=examples)
+iface = gr.Interface(predict, [gr.inputs.Image(type='filepath'),
+    gr.inputs.Radio(['style 1','style 2'], type="value", default='style 1', label='version')],
+    gr.outputs.Image(type="pil"), title=title,description=description,examples=examples)
 
 iface.launch()

model2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30a534781061f34e83bb9406b4335da4ff2616c95d22a585c1245aa8363e74e0
+size 17173511