fallenshock commited on
Commit
c46f4d1
·
verified ·
1 Parent(s): 468f99e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -9,7 +9,7 @@ import random
9
  import numpy as np
10
  import spaces
11
  import huggingface_hub
12
-
13
 
14
  from FlowEdit_utils import FlowEditSD3, FlowEditFLUX
15
  SD3STRING = 'stabilityai/stable-diffusion-3-medium-diffusers'
@@ -96,15 +96,15 @@ def FlowEditRun(
96
 
97
  if model_type == 'FLUX':
98
  # pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
99
- pipe = pipe_flux # still on CPU
100
  elif model_type == 'SD3':
101
  # pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
102
- pipe = pipe_sd3 # still on CPU
103
  else:
104
  raise NotImplementedError(f"Model type {model_type} not implemented")
105
 
106
  scheduler = pipe.scheduler
107
- pipe = pipe.to("cuda")
108
 
109
 
110
 
 
9
  import numpy as np
10
  import spaces
11
  import huggingface_hub
12
+ import copy
13
 
14
  from FlowEdit_utils import FlowEditSD3, FlowEditFLUX
15
  SD3STRING = 'stabilityai/stable-diffusion-3-medium-diffusers'
 
96
 
97
  if model_type == 'FLUX':
98
  # pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
99
+ pipe = copy.deepcopy(pipe_flux) # still on CPU
100
  elif model_type == 'SD3':
101
  # pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
102
+ pipe = copy.deepcopy(pipe_sd3) # still on CPU
103
  else:
104
  raise NotImplementedError(f"Model type {model_type} not implemented")
105
 
106
  scheduler = pipe.scheduler
107
+ pipe = pipe.to(device)
108
 
109
 
110