Update app.py

app.py

@@ -12,16 +12,18 @@ import numpy as np
 import spaces 
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+dtype = torch.bfloat16
 
 def load_models(model_path="MeissonFlow/Meissonic", 
             transformer_path="MeissonFlow/Muddit"):
     model = SymmetricTransformer2DModel.from_pretrained(
         transformer_path,
         subfolder="1024/transformer",
+        torch_dtype=dtype)
     )
-    vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae")
-    text_encoder = CLIPTextModelWithProjection.from_pretrained(model_path, subfolder="text_encoder")
-    tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+    vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae",torch_dtype=dtype)
+    text_encoder = CLIPTextModelWithProjection.from_pretrained(model_path, subfolder="text_encoder",torch_dtype=dtype)
+    tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer",torch_dtype=dtype)
     scheduler = Scheduler.from_pretrained(model_path, subfolder="scheduler")
 
     pipe = UnifiedPipeline(
@@ -65,7 +67,7 @@ def image_to_text(image, prompt, resolution=1024, steps=64, cfg=9.0):
             width=resolution,
             guidance_scale=cfg,
             num_inference_steps=steps,
-            mask_token_embedding="MeissonFlow/Muddit/1024",
+            mask_token_embedding="./mask_token_embedding.pth",
             generator=torch.manual_seed(42),
         )
 
@@ -87,7 +89,7 @@ def text_to_image(prompt, negative_prompt, num_images=1, resolution=1024, steps=
             width=resolution,
             guidance_scale=cfg,
             num_inference_steps=steps,
-            mask_token_embedding="MeissonFlow/Muddit/1024",
+            mask_token_embedding="./mask_token_embedding.pth",
             generator=torch.manual_seed(42),
         )