awacke1 commited on
Commit
7ce5fee
·
verified ·
1 Parent(s): 7691760

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -5
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(predict, [gr.inputs.Image(type='filepath'),
128
- gr.inputs.Radio(['Complex Lines','Simple Lines'], type="value", default='Simple Lines', label='version')],
129
- gr.outputs.Image(type="pil"), title=title, examples=examples)
 
 
 
 
 
 
 
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()