Update app.py

app.py

@@ -27,9 +27,11 @@ from src.pipelines.pipeline_kandinsky_subject_prior import KandinskyPriorPipelin
 from diffusers import DiffusionPipeline
 from PIL import Image
 
+__device__ = "cuda" if torch.cuda.is_available() else "cpu"
+
 class Model:
     def __init__(self):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = __device__
         
         self.text_encoder = (
             CLIPTextModelWithProjection.from_pretrained(
@@ -39,7 +41,7 @@ class Model:
             )
             .eval()
             .requires_grad_(False)
-        ).to("cuda")
+        ).to(self.device)
 
         self.tokenizer = CLIPTokenizer.from_pretrained(
             "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
@@ -99,7 +101,7 @@ class Model:
         plt.savefig('image_testt2.png')
         plt.show()
         
-        mask_img = self.image_processor(image, return_tensors="pt").to("cuda")
+        mask_img = self.image_processor(image, return_tensors="pt").to(__device__)
         vision_feats = self.vision_encoder(
             **mask_img
         ).image_embeds
@@ -111,7 +113,7 @@ class Model:
             print(indices)
         
         if image2 is not None:
-            mask_img2 = self.image_processor(image2, return_tensors="pt").to("cuda")    
+            mask_img2 = self.image_processor(image2, return_tensors="pt").to(__device__)    
             vision_feats2 = self.vision_encoder(
                 **mask_img2
             ).image_embeds
@@ -123,9 +125,9 @@ class Model:
                     print(indices)
         
         text_feats = {
-            "prompt_embeds": new_feats.text_embeds.to("cuda"),
-            "text_encoder_hidden_states": torch.tensor(new_last_hidden_states).unsqueeze(0).to("cuda"),
-            "text_mask": txt_items["attention_mask"].to("cuda"),
+            "prompt_embeds": new_feats.text_embeds.to(__device__),
+            "text_encoder_hidden_states": torch.tensor(new_last_hidden_states).unsqueeze(0).to(__device__),
+            "text_mask": txt_items["attention_mask"].to(__device__),
         }      
         return text_feats
 
@@ -219,4 +221,4 @@ def create_demo():
 
 if __name__ == '__main__':
     demo = create_demo()
-    demo.queue(max_size=20).launch()
+    demo.queue(api_open=False).launch(share=True)