Update app.py

app.py

@@ -33,7 +33,7 @@ def get_vae():
         _vae = load_vae()
     return _vae
 
-@spaces.GPU(duration=5)
+@spaces.GPU(duration=50)
 def encode_decode(img: Image.Image):
     vae = get_vae()
     img = img.convert("RGB")
@@ -45,8 +45,7 @@ def encode_decode(img: Image.Image):
     t = tfm(img).unsqueeze(0).to(device=DEVICE, dtype=DTYPE)
 
     with torch.no_grad():
-        enc = vae.encode(t)
-        lat = enc.latent_dist.mean
+        lat = vae.encode(t).latent_dist.sample()
         dec = vae.decode(lat).sample
 
     x = (dec.clamp(-1, 1) + 1) * 127.5
@@ -54,10 +53,10 @@ def encode_decode(img: Image.Image):
     out = Image.fromarray(x)
 
     # Временный PNG
-    tmp_path = os.path.join(tempfile.gettempdir(), "decoded.png")
-    out.save(tmp_path, format="PNG")
+    #tmp_path = os.path.join(tempfile.gettempdir(), "decoded.png")
+    #out.save(tmp_path, format="PNG")
 
-    return out, tmp_path
+    return out#, tmp_path
 
 with gr.Blocks(title="Asymmetric VAE 2x Upscaler") as demo:
     gr.Markdown("""
@@ -71,7 +70,7 @@ with gr.Blocks(title="Asymmetric VAE 2x Upscaler") as demo:
         out = gr.Image(type="pil", label="Decoded output")
         download = gr.File(label="Download result (PNG)")
 
-    run_btn.click(fn=encode_decode, inputs=[inp], outputs=[out, download])
+    run_btn.click(fn=encode_decode, inputs=[inp], outputs=[out]#, download])
 
 if __name__ == "__main__":
     demo.launch()