Update modules/model.py

modules/model.py

@@ -3,6 +3,7 @@ import os
 import torch
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from transformers import AutoencoderKL
+import spaces
 
 def get_checkpoints(folder):
     checkpoints = []
@@ -67,6 +68,7 @@ def get_model_and_vae_options():
 
     return formatted_models, formatted_vaes
 
+@spaces.GPU()
 def generate_image(text, neg_prompt, width, height, scheduler, num_steps, num_images, cfg_scale, seed, model, vae):
     checkpoint_folder = "../models/checkpoint/"
     vae_folder = "../models/vae/"