Update app.py

app.py

@@ -4,7 +4,7 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny
 import copy
 import random
 import time
@@ -14,8 +14,12 @@ with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 # Initialize the base model
+dtype = torch.float16
 base_model = "AlekseyCalvin/HistoricColorSoonr_Schnell"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to("cuda")
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
+torch.cuda.empty_cache()
+
 
 MAX_SEED = 2**32-1
 
@@ -56,7 +60,7 @@ def update_selection(evt: gr.SelectData, width, height):
         height,
     )
 
-@spaces.GPU(duration=70)
+@spaces.GPU(duration=50)
 def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)