skytnt commited on
Commit
72ea4ce
1 Parent(s): e94371d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -8
app.py CHANGED
@@ -5,6 +5,7 @@ import onnx
5
  import onnxruntime as rt
6
  from numpy.random import RandomState
7
  from skimage import transform
 
8
 
9
 
10
  def get_inter(r1, r2):
@@ -233,22 +234,22 @@ def gen_fn(use_seed, seed, psi):
233
  z = RandomState(int(seed) + 2 ** 31).randn(1, 512) if use_seed else np.random.randn(1, 512)
234
  w = model.get_w(z.astype(dtype=np.float32), psi)
235
  img_out = model.get_img(w)
236
- return img_out, w.tolist(), img_out
237
 
238
 
239
  def encode_img_fn(img):
240
  imgs = model.detect(img, 0.2, 0.03)
241
  if len(imgs) == 0:
242
- return "failed to detect waifu", None, None
243
  w = model.encode_img(imgs[0])
244
  img_out = model.get_img(w)
245
- return "success", imgs[0], img_out, w.tolist(), img_out
246
 
247
 
248
  def gen_video_fn(w1, w2, frame):
249
  if w1 is None or w2 is None:
250
  return None
251
- model.gen_video(np.array(w1, dtype=np.float32), np.array(w2, dtype=np.float32), "video.mp4", int(frame))
252
  return "video.mp4"
253
 
254
 
@@ -268,7 +269,7 @@ if __name__ == '__main__':
268
  gen_input3 = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.6, label="psi")
269
  gen_submit = gr.Button("Run")
270
  gen_output1 = gr.Image()
271
- select_img_input_w1 = gr.Json(visible=False)
272
  select_img_input_img1 = gr.Image(visible=False)
273
 
274
  with gr.TabItem("encode image"):
@@ -279,7 +280,7 @@ if __name__ == '__main__':
279
  with gr.Row():
280
  encode_img_output2 = gr.Image(label="detected")
281
  encode_img_output3 = gr.Image(label="encoded")
282
- select_img_input_w2 = gr.Json(visible=False)
283
  select_img_input_img2 = gr.Image(visible=False)
284
 
285
  with gr.TabItem("generate video"):
@@ -293,7 +294,7 @@ if __name__ == '__main__':
293
  "current encoded image"], type="index")
294
  select_img1_button = gr.Button("select")
295
  select_img1_output_img = gr.Image(label="image 1")
296
- select_img1_output_w = gr.Json(visible=False)
297
  with gr.Column():
298
  gr.Markdown("please select image 2")
299
  select_img2_dropdown = gr.Dropdown(label="source",
@@ -301,7 +302,7 @@ if __name__ == '__main__':
301
  "current encoded image"], type="index")
302
  select_img2_button = gr.Button("select")
303
  select_img2_output_img = gr.Image(label="image 2")
304
- select_img2_output_w = gr.Json(visible=False)
305
  generate_video_frame = gr.Slider(minimum=10, maximum=30, step=1, label="frame", value=10)
306
  generate_video_button = gr.Button("generate")
307
  generate_video_output = gr.Video()
 
5
  import onnxruntime as rt
6
  from numpy.random import RandomState
7
  from skimage import transform
8
+ import json
9
 
10
 
11
  def get_inter(r1, r2):
 
234
  z = RandomState(int(seed) + 2 ** 31).randn(1, 512) if use_seed else np.random.randn(1, 512)
235
  w = model.get_w(z.astype(dtype=np.float32), psi)
236
  img_out = model.get_img(w)
237
+ return img_out, json.dumps(w.tolist()), img_out
238
 
239
 
240
  def encode_img_fn(img):
241
  imgs = model.detect(img, 0.2, 0.03)
242
  if len(imgs) == 0:
243
+ return "failed to detect waifu", None, None, None, None
244
  w = model.encode_img(imgs[0])
245
  img_out = model.get_img(w)
246
+ return "success", imgs[0], img_out, json.dumps(w.tolist()), img_out
247
 
248
 
249
  def gen_video_fn(w1, w2, frame):
250
  if w1 is None or w2 is None:
251
  return None
252
+ model.gen_video(np.array(json.loads(w1), dtype=np.float32), np.array(json.loads(w2), dtype=np.float32), "video.mp4", int(frame))
253
  return "video.mp4"
254
 
255
 
 
269
  gen_input3 = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.6, label="psi")
270
  gen_submit = gr.Button("Run")
271
  gen_output1 = gr.Image()
272
+ select_img_input_w1 = gr.Textbox(visible=False)
273
  select_img_input_img1 = gr.Image(visible=False)
274
 
275
  with gr.TabItem("encode image"):
 
280
  with gr.Row():
281
  encode_img_output2 = gr.Image(label="detected")
282
  encode_img_output3 = gr.Image(label="encoded")
283
+ select_img_input_w2 = gr.Textbox(visible=False)
284
  select_img_input_img2 = gr.Image(visible=False)
285
 
286
  with gr.TabItem("generate video"):
 
294
  "current encoded image"], type="index")
295
  select_img1_button = gr.Button("select")
296
  select_img1_output_img = gr.Image(label="image 1")
297
+ select_img1_output_w = gr.Textbox(visible=False)
298
  with gr.Column():
299
  gr.Markdown("please select image 2")
300
  select_img2_dropdown = gr.Dropdown(label="source",
 
302
  "current encoded image"], type="index")
303
  select_img2_button = gr.Button("select")
304
  select_img2_output_img = gr.Image(label="image 2")
305
+ select_img2_output_w = gr.Textbox(visible=False)
306
  generate_video_frame = gr.Slider(minimum=10, maximum=30, step=1, label="frame", value=10)
307
  generate_video_button = gr.Button("generate")
308
  generate_video_output = gr.Video()