add gpu wrapper

app.py

@@ -10,6 +10,7 @@ import matplotlib.pyplot as plt
 from matplotlib.colors import hex2color
 import stone
 import os
+import spaces
 
 # Define model initialization functions
 def load_model(model_name):
@@ -47,27 +48,29 @@ def load_model(model_name):
         raise ValueError("Unknown model name")
     return pipeline
 
-# choices=[
-#             "stabilityai/sdxl-turbo", 
-#             "runwayml/stable-diffusion-v1-5", 
-#             "ByteDance/SDXL-Lightning", 
-#             "segmind/SSD-1B"
-#         ]
+choices=[
+            "stabilityai/sdxl-turbo", 
+            "runwayml/stable-diffusion-v1-5", 
+            "ByteDance/SDXL-Lightning", 
+            "segmind/SSD-1B"
+        ]
 
-# for model_name in choices:
-#     load_model(model_name)
+for model_name in choices:
+    load_model(model_name)
 
 # Initialize the default model
 default_model = "stabilityai/sdxl-turbo"
 
 pipeline_text2image = load_model(default_model)
 
+@spaces.GPU
 def getimgen(prompt):
     return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
 
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
 
+@spaces.GPU
 def blip_caption_image(image, prefix):
     inputs = blip_processor(image, prefix, return_tensors="pt").to("cuda", torch.float16)
     out = blip_model.generate(**inputs)
@@ -106,6 +109,7 @@ def skintoneplot(hex_codes):
             ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
     return fig
 
+@spaces.GPU
 def generate_images_plots(prompt, model_name):
     global pipeline_text2image
     pipeline_text2image = load_model(model_name)