Update app.py

app.py

@@ -9,8 +9,7 @@ import model_build
 
 
 psp = model_build.build_psp()
-stylegan2_v1 = model_build.build_stylegan2(network_pkl = 'pretrained/ohayou_face.pkl')
-stylegan2_v2 = model_build.build_stylegan2(network_pkl = 'pretrained/ohayou_face2.pkl')
+stylegan2 = model_build.build_stylegan2(network_pkl = 'pretrained/ohayou_face2.pkl')
 
 pretransform = transforms.Compose([
     transforms.Resize((256, 256)),
@@ -21,13 +20,12 @@ def pipeline(img, v):
     img = model_build.img_preprocess(img, pretransform)
     with torch.no_grad():
         latent_space = psp(img.float(), randomize_noise=True, resize=False, return_latents=True)
-    if v == 'version1': img = stylegan2_v1(latent_space, noise_mode='none')
-    elif v == 'version2': img = stylegan2_v2(latent_space, noise_mode='none')
+    img = stylegan2(latent_space, noise_mode='none')
     img = Image.fromarray(np.array((img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).squeeze(0)[20:-20,:,:]))
     img.save('output.png')
     return 'output.png'
     
 examples=[['1.png',False], ['2.png', False], ['3.png', False], ['4.png', False], ['5.png', False]]
 description="이미지 업로드 후 연필모양 아이콘 누르면 정사각형으로 자르기 할 수 있습니다. 디테일은 좀 떨어짐. 잘 안되면 얼굴 클로즈업하면 좀 나아요. After uploading the image, you can square crop."
-gr.Interface(pipeline, [gr.inputs.Image(type="pil"), gr.inputs.Radio(["version1", "version2"])], gr.outputs.Image(type="file"),description=description,allow_flagging=False,examples=examples,allow_screenshot=False,enable_queue=False).launch()
+gr.Interface(pipeline, [gr.inputs.Image(type="pil")], gr.outputs.Image(type="file"),description=description,allow_flagging=False,examples=examples,allow_screenshot=False,enable_queue=False).launch()