Jordan Legg commited on
Commit
d83ac73
·
1 Parent(s): 2c3bf55

test xformers

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -8,11 +8,13 @@ from typing import Tuple, Any
8
 
9
  dtype: torch.dtype = torch.bfloat16
10
  device: str = "cuda" if torch.cuda.is_available() else "cpu"
11
- MAX_SEED = np.iinfo(np.int64).max
12
 
13
  pipe = DiffusionPipeline.from_pretrained("shuttleai/shuttle-3-diffusion", torch_dtype=dtype).to(device)
14
  # Enable VAE tiling
15
  pipe.vae.enable_tiling()
 
 
16
 
17
  # Define cinematic aspect ratios
18
  ASPECT_RATIOS = {
@@ -29,7 +31,7 @@ ASPECT_RATIOS = {
29
  "1.43:1 (IMAX 70mm)": 1.43,
30
  "2.40:1 (Modern Anamorphic)": 2.40
31
  }
32
-
33
  MIN_WIDTH = 512
34
  MAX_WIDTH = 3072
35
  STANDARD_WIDTH = 2048
 
8
 
9
  dtype: torch.dtype = torch.bfloat16
10
  device: str = "cuda" if torch.cuda.is_available() else "cpu"
11
+
12
 
13
  pipe = DiffusionPipeline.from_pretrained("shuttleai/shuttle-3-diffusion", torch_dtype=dtype).to(device)
14
  # Enable VAE tiling
15
  pipe.vae.enable_tiling()
16
+ # Enable memory efficient attention
17
+ pipe.enable_xformers_memory_efficient_attention()
18
 
19
  # Define cinematic aspect ratios
20
  ASPECT_RATIOS = {
 
31
  "1.43:1 (IMAX 70mm)": 1.43,
32
  "2.40:1 (Modern Anamorphic)": 2.40
33
  }
34
+ MAX_SEED = np.iinfo(np.int64).max
35
  MIN_WIDTH = 512
36
  MAX_WIDTH = 3072
37
  STANDARD_WIDTH = 2048