Spaces:
Running
on
Zero
Running
on
Zero
Prime Cai
commited on
Commit
·
a1ef85e
1
Parent(s):
f3e9f34
dynamic gpu
Browse files
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 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
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)
|