Update app.py
Browse files
app.py
CHANGED
@@ -88,11 +88,11 @@ class Generator(nn.Module):
|
|
88 |
|
89 |
# 🔧 Load the models
|
90 |
model1 = Generator(3, 1, 3)
|
91 |
-
model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
|
92 |
model1.eval()
|
93 |
|
94 |
model2 = Generator(3, 1, 3)
|
95 |
-
model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu')))
|
96 |
model2.eval()
|
97 |
|
98 |
# 🖼️ Function to process the image and create line drawing
|
@@ -124,8 +124,15 @@ for file in os.listdir(image_dir):
|
|
124 |
examples.append([file, 'Complex Lines'])
|
125 |
|
126 |
# 🚀 Create and launch the Gradio interface
|
127 |
-
iface = gr.Interface(
|
128 |
-
|
129 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
130 |
|
131 |
iface.launch()
|
|
|
88 |
|
89 |
# 🔧 Load the models
|
90 |
model1 = Generator(3, 1, 3)
|
91 |
+
model1.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu'), weights_only=True))
|
92 |
model1.eval()
|
93 |
|
94 |
model2 = Generator(3, 1, 3)
|
95 |
+
model2.load_state_dict(torch.load('model2.pth', map_location=torch.device('cpu'), weights_only=True))
|
96 |
model2.eval()
|
97 |
|
98 |
# 🖼️ Function to process the image and create line drawing
|
|
|
124 |
examples.append([file, 'Complex Lines'])
|
125 |
|
126 |
# 🚀 Create and launch the Gradio interface
|
127 |
+
iface = gr.Interface(
|
128 |
+
fn=predict,
|
129 |
+
inputs=[
|
130 |
+
gr.Image(type='filepath'),
|
131 |
+
gr.Radio(['Complex Lines', 'Simple Lines'], label='version', value='Simple Lines')
|
132 |
+
],
|
133 |
+
outputs=gr.Image(type="pil"),
|
134 |
+
title=title,
|
135 |
+
examples=examples
|
136 |
+
)
|
137 |
|
138 |
iface.launch()
|