Prime Cai commited on
Commit
a1ef85e
·
1 Parent(s): f3e9f34

dynamic gpu

Browse files
Files changed (1) hide show
  1. app.py +20 -27
app.py CHANGED
@@ -16,29 +16,26 @@ CHECKPOINT = "primecai/dsd_model"
16
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
  dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
18
 
19
- @spaces.GPU
20
- def init_pipeline():
21
- global pipe
22
- transformer = FluxTransformer2DConditionalModel.from_pretrained(
23
- CHECKPOINT,
24
- subfolder="transformer",
25
- torch_dtype=dtype,
26
- low_cpu_mem_usage=False,
27
- ignore_mismatched_sizes=True,
28
- use_auth_token=os.getenv("HF_TOKEN"),
29
- )
30
- pipe = FluxConditionalPipeline.from_pretrained(
31
- "black-forest-labs/FLUX.1-dev",
32
- transformer=transformer,
33
- torch_dtype=dtype,
34
- use_auth_token=os.getenv("HF_TOKEN"),
35
- )
36
- pipe.load_lora_weights(
37
- CHECKPOINT,
38
- weight_name="pytorch_lora_weights.safetensors",
39
- use_auth_token=os.getenv("HF_TOKEN"),
40
- )
41
- pipe.to(device, dtype=dtype)
42
 
43
  @spaces.GPU
44
  def generate_image(image, text, gemini_prompt, guidance, i_guidance, t_guidance):
@@ -47,9 +44,6 @@ def generate_image(image, text, gemini_prompt, guidance, i_guidance, t_guidance)
47
  ((w - min_size) // 2, (h - min_size) // 2, (w + min_size) // 2, (h + min_size) // 2)
48
  ).resize((512, 512))
49
 
50
- if pipe is None:
51
- init_pipeline()
52
-
53
  control_image = load_image(image)
54
  result_image = pipe(
55
  prompt=text.strip(),
@@ -134,5 +128,4 @@ with demo:
134
  )
135
 
136
  if __name__ == "__main__":
137
- init_pipeline()
138
  demo.launch(debug=False, share=True)
 
16
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
  dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
18
 
19
+ transformer = FluxTransformer2DConditionalModel.from_pretrained(
20
+ CHECKPOINT,
21
+ subfolder="transformer",
22
+ torch_dtype=dtype,
23
+ low_cpu_mem_usage=False,
24
+ ignore_mismatched_sizes=True,
25
+ use_auth_token=os.getenv("HF_TOKEN"),
26
+ )
27
+ pipe = FluxConditionalPipeline.from_pretrained(
28
+ "black-forest-labs/FLUX.1-dev",
29
+ transformer=transformer,
30
+ torch_dtype=dtype,
31
+ use_auth_token=os.getenv("HF_TOKEN"),
32
+ )
33
+ pipe.load_lora_weights(
34
+ CHECKPOINT,
35
+ weight_name="pytorch_lora_weights.safetensors",
36
+ use_auth_token=os.getenv("HF_TOKEN"),
37
+ )
38
+ pipe.to(device, dtype=dtype)
 
 
 
39
 
40
  @spaces.GPU
41
  def generate_image(image, text, gemini_prompt, guidance, i_guidance, t_guidance):
 
44
  ((w - min_size) // 2, (h - min_size) // 2, (w + min_size) // 2, (h + min_size) // 2)
45
  ).resize((512, 512))
46
 
 
 
 
47
  control_image = load_image(image)
48
  result_image = pipe(
49
  prompt=text.strip(),
 
128
  )
129
 
130
  if __name__ == "__main__":
 
131
  demo.launch(debug=False, share=True)