#15931: Re-enable SD nightly and demo tests

tenstorrent · Jan 23, 2025 · b0986ea · b0986ea
1 parent 182c424
commit b0986ea
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 12 deletions.
diff --git a/models/demos/wormhole/stable_diffusion/demo/demo.py b/models/demos/wormhole/stable_diffusion/demo/demo.py
@@ -18,7 +18,6 @@
 from diffusers import (
     AutoencoderKL,
     UNet2DConditionModel,
-    StableDiffusionPipeline,
 )
 from models.utility_functions import skip_for_grayskull
 from models.utility_functions import (
@@ -31,11 +30,9 @@
 from models.demos.wormhole.stable_diffusion.tt.ttnn_functional_unet_2d_condition_model_new_conv import (
     UNet2DConditionModel as UNet2D,
 )
-from models.demos.wormhole.stable_diffusion.tt.ttnn_functional_utility_functions import round_up_to_tile_dim
 from torchvision.transforms import ToTensor
 from torchmetrics.multimodal.clip_score import CLIPScore
 from torchmetrics.image.fid import FrechetInceptionDistance
-from scipy import integrate
 
 
 def load_inputs(input_path):
@@ -598,7 +595,6 @@ def run_demo_inference_diffusiondb(
         logger.info(f"CLIP Score (TTNN): {clip_score_ttnn}")
 
 
-@pytest.mark.skip(reason="#9945: Skip for now since this breaks on WH because of di/dt")
 @skip_for_grayskull()
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
 @pytest.mark.parametrize(

diff --git a/models/demos/wormhole/stable_diffusion/tests/test_unet_2d_condition_model.py b/models/demos/wormhole/stable_diffusion/tests/test_unet_2d_condition_model.py
@@ -11,9 +11,7 @@
 
 from models.utility_functions import (
     skip_for_grayskull,
-    is_wormhole_b0,
     comp_pcc,
-    is_blackhole,
 )
 from diffusers import LMSDiscreteScheduler
 import ttnn
@@ -72,10 +70,11 @@ def unsqueeze_all_params_to_4d(params):
         (2, 4, 64, 64),
     ],
 )
-@pytest.mark.skip(reason="#15931: Failing, skip for now")
 def test_unet_2d_condition_model_512x512(device, batch_size, in_channels, input_height, input_width):
     device.enable_program_cache()
 
+    os.environ["SLOW_MATMULS"] = "1"
+
     # setup envvar if testing on N300
     wh_arch_yaml_org = None
     if device.core_grid.y == 7:

diff --git a/models/demos/wormhole/stable_diffusion/tests/test_upblock_2d.py b/models/demos/wormhole/stable_diffusion/tests/test_upblock_2d.py
@@ -5,10 +5,11 @@
 import torch
 from diffusers import StableDiffusionPipeline
 from loguru import logger
+import os
 import ttnn
 import pytest
 
-from models.utility_functions import tt_to_torch_tensor, torch_random
+from models.utility_functions import torch_random
 from tests.ttnn.utils_for_testing import assert_with_pcc
 from models.utility_functions import (
     skip_for_grayskull,
@@ -18,7 +19,6 @@
 from models.demos.wormhole.stable_diffusion.custom_preprocessing import custom_preprocessor
 from ttnn.model_preprocessing import preprocess_model_parameters
 from models.demos.wormhole.stable_diffusion.tt.ttnn_functional_utility_functions import (
-    pre_process_input,
     post_process_output,
     weight_to_bfp8,
 )
@@ -29,10 +29,9 @@
 @pytest.mark.parametrize("res_hidden_states_tuple", [([2, 1280, 8, 8], [2, 1280, 8, 8], [2, 1280, 8, 8])])
 @pytest.mark.parametrize("hidden_states", [[2, 1280, 8, 8]])
 @pytest.mark.parametrize("temb", [[1, 1, 2, 1280]])
-@pytest.mark.skip(reason="#15931: Fails, need to investigate")
 def test_upblock_512x512(reset_seeds, device, res_hidden_states_tuple, hidden_states, temb):
-    # TODO
-    # setup pytorch model
+    os.environ["SLOW_MATMULS"] = "1"
+
     pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
     unet = pipe.unet
     unet.eval()