Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 
 import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline ,AutoencoderTiny
+#from diffusers import DiffusionPipeline ,AutoencoderTiny
 import torch
 
 from huggingface_hub import login
@@ -13,10 +13,10 @@ login(token=a )
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 #model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-model_repo_id = "stabilityai/stable-diffusion-xl-base-0.9"
-
+#model_repo_id = "stabilityai/stable-diffusion-xl-base-0.9"
 
 
+"""
 if torch.cuda.is_available():
     torch_dtype = torch.float16
 else:
@@ -24,11 +24,21 @@ else:
 
     pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
+"""
+
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
+from diffusers import AutoencoderTiny, StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "nota-ai/bk-sdm-small", torch_dtype=torch.float16, use_safetensors=True)
+pipe.vae = AutoencoderTiny.from_pretrained(
+    "sayakpaul/taesd-diffusers", torch_dtype=torch.float16, use_safetensors=True)
+pipe = pipe.to(device)
+
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,