Spaces:
vilarin
/
Running on Zero

vilarin commited on
Commit
5d300f8
1 Parent(s): ca39da7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -1
app.py CHANGED
@@ -34,12 +34,19 @@ JS = """function () {
34
  }"""
35
 
36
 
 
 
 
 
 
 
37
 
38
 
39
  # Ensure model and scheduler are initialized in GPU-enabled function
40
  if torch.cuda.is_available():
41
- pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to("cuda")
42
 
 
43
 
44
  # Function
45
  @spaces.GPU()
 
34
  }"""
35
 
36
 
37
+ vae = AutoencoderKL.from_pretrained(
38
+ repo,
39
+ subfolder="vae",
40
+ torch_dtype=torch.float16,
41
+ )
42
+
43
 
44
 
45
  # Ensure model and scheduler are initialized in GPU-enabled function
46
  if torch.cuda.is_available():
47
+ pipe = StableDiffusion3Pipeline.from_pretrained(repo, vae=vae, torch_dtype=torch.float16).to("cuda")
48
 
49
+ pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
50
 
51
  # Function
52
  @spaces.GPU()