changes for ZeroGPU

app.py

@@ -31,11 +31,13 @@ from src.priors.prior_transformer import (
 from src.pipelines.pipeline_kandinsky_prior import KandinskyPriorPipeline
 
 from diffusers import DiffusionPipeline
+import spaces
 
 
 __DEVICE__ = "cpu"
 if torch.cuda.is_available():
     __DEVICE__ = "cuda"
+__DEVICE__ = "cuda"
 
 class Ours:
     def __init__(self, device):
@@ -43,7 +45,7 @@ class Ours:
             CLIPTextModelWithProjection.from_pretrained(
                 "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
                 projection_dim=1280,
-                torch_dtype=torch.float32,
+                torch_dtype=torch.float16,
             )
             .eval()
             .requires_grad_(False)
@@ -55,7 +57,7 @@ class Ours:
 
         prior = PriorTransformer.from_pretrained(
             "ECLIPSE-Community/ECLIPSE_KandinskyV22_Prior",
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
         )
 
         self.pipe_prior = KandinskyPriorPipeline.from_pretrained(
@@ -63,16 +65,16 @@ class Ours:
             prior=prior,
             text_encoder=text_encoder,
             tokenizer=tokenizer,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
         ).to(device)
 
         self.pipe = DiffusionPipeline.from_pretrained(
-            "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float32
+            "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
         ).to(device)
 
     def inference(self, text, negative_text, steps, guidance_scale):
         gen_images = []
-        for i in range(1):
+        for i in range(4):
             image_emb, negative_image_emb = self.pipe_prior(
                 text, negative_prompt=negative_text
             ).to_tuple()
@@ -88,13 +90,13 @@ class Ours:
 
 selected_model = Ours(device=__DEVICE__)
 
-
+@spaces.GPU
 def get_images(text, negative_text, steps, guidance_scale):
     images = selected_model.inference(text, negative_text, steps, guidance_scale)
     new_images = []
     for img in images:
         new_images.append(img)
-    return new_images[0]
+    return new_images
 
 
 with gr.Blocks() as demo:
@@ -137,9 +139,9 @@ with gr.Blocks() as demo:
         with gr.Row():
             btn = gr.Button(value="Generate Image")
 
-    gallery = gr.Image(
-        height=512, width=512, label="Generated images", show_label=True, elem_id="gallery"
-    )
+    gallery = gr.Gallery(
+        label="Generated images", show_label=False, elem_id="gallery"
+    , columns=[4], rows=[1], object_fit="contain", height="auto")
 
     btn.click(
         get_images,