clean code

pipeline_freescale.py

@@ -33,6 +33,7 @@ from inspect import isfunction
 from functools import partial
 import numpy as np
 
+import torch.nn.functional as F
 from diffusers.models.attention import BasicTransformerBlock
 from scale_attention import ori_forward, scale_forward
 

pipeline_freescale_turbo.py

@@ -33,6 +33,7 @@ from inspect import isfunction
 from functools import partial
 import numpy as np
 
+import torch.nn.functional as F
 from diffusers.models.attention import BasicTransformerBlock
 from scale_attention_turbo import ori_forward, scale_forward
 
@@ -815,7 +816,7 @@ class StableDiffusionXLPipeline_Turbo(DiffusionPipeline, FromSingleFileMixin, Lo
             height, width = resolutions_list[0]
             target_sizes = resolutions_list[1:]
             if not restart_steps:
-                restart_steps = [1] * len(target_sizes)
+                restart_steps = [int(num_inference_steps*0.5)] * len(target_sizes)
         else:
             height = height or self.default_sample_size * self.vae_scale_factor
             width = width or self.default_sample_size * self.vae_scale_factor

scale_attention.py

@@ -2,8 +2,7 @@ from typing import Any, Dict, Optional
 
 import torch
 import torch.nn.functional as F
-from torch import nn
-from einops import rearrange, repeat
+from einops import rearrange
 import random
 
 def gaussian_kernel(kernel_size=3, sigma=1.0, channels=3):

scale_attention_turbo.py

@@ -2,8 +2,7 @@ from typing import Any, Dict, Optional
 
 import torch
 import torch.nn.functional as F
-from torch import nn
-from einops import rearrange, repeat
+from einops import rearrange
 import random
 
 def gaussian_kernel(kernel_size=3, sigma=1.0, channels=3):