add gpu wrapper

app.py

@@ -13,6 +13,7 @@ import os
 import spaces
 
 # Define model initialization functions
+@spaces.GPU
 def load_model(model_name):
     if model_name == "stabilityai/sdxl-turbo":
         pipeline = DiffusionPipeline.from_pretrained(
@@ -48,24 +49,21 @@ 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"
-        ]
-
-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]
+def getimgen(prompt, model_name):
+    if model_name == "stabilityai/sdxl-turbo":
+        return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
+    elif model_name == "runwayml/stable-diffusion-v1-5":
+        return pipeline_text2image(prompt).images[0]
+    elif model_name == "ByteDance/SDXL-Lightning":
+        return pipeline_text2image(prompt, num_inference_steps=4, guidance_scale=0).images[0]
+    elif model_name == "segmind/SSD-1B":
+        neg_prompt = "ugly, blurry, poor quality"
+        return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt).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")
@@ -115,7 +113,7 @@ def generate_images_plots(prompt, model_name):
     pipeline_text2image = load_model(model_name)
     foldername = "temp"
     Path(foldername).mkdir(parents=True, exist_ok=True)
-    images = [getimgen(prompt) for _ in range(10)]
+    images = [getimgen(prompt, model_name) for _ in range(10)]
     genders = []
     skintones = []
     for image, i in zip(images, range(10)):
@@ -159,4 +157,4 @@ with gr.Blocks(title="Skin Tone and Gender bias in Text to Image Models") as dem
         genplot = gr.Plot(label="Gender")
     btn.click(generate_images_plots, inputs=[prompt, model_dropdown], outputs=[gallery, skinplot, genplot])
 
-demo.launch(debug=True)
+demo.launch(debug=True)