Update app.py

app.py

@@ -16,11 +16,6 @@ from flextok.flextok_wrapper import FlexTokFromHub
 from flextok.utils.demo import imgs_from_urls, denormalize, batch_to_pil
 from flextok.utils.misc import detect_bf16_support, get_bf16_context, get_generator
 
-#import torch._dynamo
-#torch._dynamo.config.suppress_errors = True
-os.environ["TORCH_LOGS"] = "+dynamo"
-os.environ["TORCHDYNAMO_VERBOSE"] = "1"
-
 # We recommend running this demo on an A100 GPU
 if torch.cuda.is_available():
     device = "cuda"
@@ -53,8 +48,14 @@ MODEL_NAME = 'FlexTok d18-d28 (DFN)'
 
 # Load FlexTok model from HF Hub
 flextok_model = FlexTokFromHub.from_pretrained(MODEL_ID).to(device).eval()
-#flextok_model.encoder.module_dict.enc_seq_packer.compile_block_mask = False
-#flextok_model.decoder.module_dict.dec_seq_packer.compile_block_mask = False
+
+# Disable flex_attention for HF Space
+flextok_model.encoder.module_dict.enc_seq_packer.return_materialized_mask = True
+flextok_model.decoder.module_dict.dec_seq_packer.return_materialized_mask = True
+for block in flextok_model.encoder.module_dict.enc_transformer.blocks:
+    block._checkpoint_wrapped_module.attn.use_flex_attention = False
+for block in flextok_model.decoder.module_dict.dec_transformer.blocks:
+    block._checkpoint_wrapped_module.attn.use_flex_attention = False
 
 # Load AuraSR model from HF Hub
 try: