Update app.py

app.py

@@ -1,37 +1,50 @@
-from pathlib import Path
-import torch
-import spaces
-import shutil
 import gradio as gr
-from loguru import logger
 import requests
-from tqdm import tqdm
-from diffusers import AutoPipelineForImage2Image
+import shutil
+import spaces
+import torch
+from diffusers import AutoencoderKL, StableDiffusionXLImg2ImgPipeline
+from loguru import logger
+from pathlib import Path
 from PIL import Image
+from tqdm import tqdm
 
-model_path = Path("pony-diffusion-v6-xl.safetensors")
 
-if not model_path.exists():
-    logger.info(f"Downloading model: {model_path}")
+def download(file: str, url: str):
+    file_path = Path(file)
+    if file_path.exists():
+        return
 
-    r = requests.get(
-        "https://civitai.com/api/download/models/290640?type=Model&format=SafeTensor&size=pruned&fp=fp16",
-        stream=True,
-    )
+    r = requests.get(url, stream=True)
     r.raise_for_status()
 
-    temp_path = f"/tmp/{model_path.name}"
+    temp_path = f"/tmp/{file_path.name}"
     with tqdm(
-        total=int(r.headers["content-length"]), unit="B", unit_scale=True
+        desc=file, total=int(r.headers["content-length"]), unit="B", unit_scale=True
     ) as pbar, open(temp_path, "wb") as f:
         for chunk in r.iter_content(chunk_size=1024 * 1024):
             f.write(chunk)
             pbar.update(len(chunk))
 
-    shutil.move(temp_path, model_path)
+    shutil.move(temp_path, file_path)
+
+
+model_path = "pony-diffusion-v6-xl.safetensors"
+download(
+    model_path,
+    "https://civitai.com/api/download/models/290640?type=Model&format=SafeTensor&size=pruned&fp=fp16",
+)
+
+vae_path = "pony-diffusion-v6-xl.vae.safetensors"
+download(
+    vae_path,
+    "https://civitai.com/api/download/models/290640?type=VAE&format=SafeTensor",
+)
+
+vae = AutoencoderKL.from_single_file(vae_path)
 
-pipe = AutoPipelineForImage2Image.from_pretrained(
-    model_path, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
+pipe = StableDiffusionXLImg2ImgPipeline.from_single_file(
+    model_path, torch_dtype=torch.float16, use_safetensors=True, variant="fp16", vae=vae
 )
 pipe = pipe.to("cuda")