Update app.py

app.py

@@ -4,73 +4,120 @@ import torchvision.transforms as T
 from PIL import Image
 from diffusers import AsymmetricAutoencoderKL
 import spaces
-import io
-import tempfile
-import os
 
 MODEL_ID = "babkasotona/vae8x16x32ch"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.float16
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
-def load_vae(model_id=MODEL_ID, device=DEVICE):
+
+# -------------------------
+# Load VAE
+# -------------------------
+def load_vae(model_id=MODEL_ID):
     for attempt in (None, "vae"):
         try:
             if attempt is None:
-                vae = AsymmetricAutoencoderKL.from_pretrained(model_id, torch_dtype=DTYPE)
+                vae = AsymmetricAutoencoderKL.from_pretrained(
+                    model_id,
+                    torch_dtype=DTYPE
+                )
             else:
-                vae = AsymmetricAutoencoderKL.from_pretrained(model_id, subfolder=attempt, torch_dtype=DTYPE)
-            vae.to(device)
-            vae.eval().half()
+                vae = AsymmetricAutoencoderKL.from_pretrained(
+                    model_id,
+                    subfolder=attempt,
+                    torch_dtype=DTYPE
+                )
+
+            vae = vae.to(DEVICE)
+            vae.eval()
+
+            print("VAE loaded on", DEVICE)
             return vae
+
         except Exception as e:
             last_err = e
-    raise RuntimeError(f"Failed to load VAE {model_id}: {last_err}")
+
+    raise RuntimeError(f"Failed to load VAE: {last_err}")
+
 
 _vae = None
+
+
 def get_vae():
     global _vae
     if _vae is None:
         _vae = load_vae()
     return _vae
 
+
+# -------------------------
+# Encode / Decode
+# -------------------------
 @spaces.GPU(duration=50)
 def encode_decode(img: Image.Image):
+
+    if img is None:
+        raise gr.Error("Please upload an image")
+
     vae = get_vae()
+
     img = img.convert("RGB")
 
     tfm = T.Compose([
         T.ToTensor(),
-        T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        T.Normalize([0.5]*3, [0.5]*3),
     ])
-    t = tfm(img).unsqueeze(0).to(device=DEVICE, dtype=DTYPE)
+
+    t = tfm(img).unsqueeze(0).to(DEVICE, dtype=DTYPE)
+
+    print("Input tensor:", t.shape, t.dtype, t.device)
 
     with torch.no_grad():
-        lat = vae.encode(t).latent_dist.sample()
+
+        enc = vae.encode(t)
+        lat = enc.latent_dist.sample()
+
+        print("Latents:", lat.shape)
+
         dec = vae.decode(lat).sample
 
     x = (dec.clamp(-1, 1) + 1) * 127.5
-    x = x.round().to(torch.uint8).squeeze(0).permute(1, 2, 0).cpu().numpy()
+    x = x.round().to(torch.uint8)
+
+    x = x.squeeze(0).permute(1, 2, 0).cpu().numpy()
+
     out = Image.fromarray(x)
 
-    # Временный PNG
-    #tmp_path = os.path.join(tempfile.gettempdir(), "decoded.png")
-    #out.save(tmp_path, format="PNG")
+    print("Output size:", out.size)
+
+    return out
 
-    return out#, tmp_path
 
+# -------------------------
+# UI
+# -------------------------
 with gr.Blocks(title="Asymmetric VAE 2x Upscaler") as demo:
-    gr.Markdown("""
-    # 🧠 Asymmetric VAE 2x Upscaler  
-    Загрузите изображение → нажмите **"Upscale"** 
-    """)
 
-    with gr.Column():
-        inp = gr.Image(type="pil", label="Upload image")
-        run_btn = gr.Button("Upscale")
-        out = gr.Image(type="pil", label="Decoded output")
-        download = gr.File(label="Download result (PNG)")
+    gr.Markdown(
+        "# 🧠 Asymmetric VAE Upscaler\n"
+        "Upload image → press **Upscale**"
+    )
+
+    inp = gr.Image(type="pil", label="Upload image")
+
+    run_btn = gr.Button("Upscale")
+
+    out = gr.Image(label="Decoded output")
+
+    run_btn.click(
+        fn=encode_decode,
+        inputs=inp,
+        outputs=out
+    )
 
-    run_btn.click(fn=encode_decode, inputs=[inp], outputs=[out])#, download])
 
+# -------------------------
+# Launch
+# -------------------------
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()