Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
|
@@ -12,16 +12,16 @@ pipe_14b = Cosmos2TextToImagePipeline.from_pretrained(model_14b_id,
|
|
| 12 |
torch_dtype=torch.bfloat16
|
| 13 |
)
|
| 14 |
|
| 15 |
-
pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
|
| 16 |
-
revision="refs/pr/2",
|
| 17 |
-
text_encoder=pipe_14b.text_encoder,
|
| 18 |
-
safety_checker=pipe_14b.safety_checker,
|
| 19 |
-
vae=pipe_14b.vae,
|
| 20 |
-
tokenizer=pipe_14b.tokenizer,
|
| 21 |
-
torch_dtype=torch.bfloat16
|
| 22 |
)
|
| 23 |
pipe_14b.to("cuda")
|
| 24 |
-
pipe_2b.to("cuda")
|
| 25 |
|
| 26 |
scheduler = EDMEulerScheduler(
|
| 27 |
sigma_min=0.002,
|
|
@@ -47,16 +47,13 @@ def generate_image(prompt,
|
|
| 47 |
progress=gr.Progress(track_tqdm=True)
|
| 48 |
):
|
| 49 |
|
| 50 |
-
pipe_14b.to("cuda")
|
| 51 |
-
pipe_2b.to("cuda")
|
| 52 |
-
|
| 53 |
if model_choice == "14B":
|
| 54 |
pipe = pipe_14b
|
| 55 |
-
pipe_2b.to("cpu")
|
| 56 |
|
| 57 |
-
else:
|
| 58 |
-
|
| 59 |
-
|
| 60 |
|
| 61 |
if randomize_seed:
|
| 62 |
actual_seed = random.randint(0, 1000000)
|
|
|
|
| 12 |
torch_dtype=torch.bfloat16
|
| 13 |
)
|
| 14 |
|
| 15 |
+
#pipe_2b = Cosmos2TextToImagePipeline.from_pretrained(model_2b_id,
|
| 16 |
+
# revision="refs/pr/2",
|
| 17 |
+
# text_encoder=pipe_14b.text_encoder,
|
| 18 |
+
# safety_checker=pipe_14b.safety_checker,
|
| 19 |
+
# vae=pipe_14b.vae,
|
| 20 |
+
# tokenizer=pipe_14b.tokenizer,
|
| 21 |
+
# torch_dtype=torch.bfloat16
|
| 22 |
)
|
| 23 |
pipe_14b.to("cuda")
|
| 24 |
+
#pipe_2b.to("cuda")
|
| 25 |
|
| 26 |
scheduler = EDMEulerScheduler(
|
| 27 |
sigma_min=0.002,
|
|
|
|
| 47 |
progress=gr.Progress(track_tqdm=True)
|
| 48 |
):
|
| 49 |
|
|
|
|
|
|
|
|
|
|
| 50 |
if model_choice == "14B":
|
| 51 |
pipe = pipe_14b
|
| 52 |
+
#pipe_2b.to("cpu")
|
| 53 |
|
| 54 |
+
#else:
|
| 55 |
+
# pipe = pipe_2b
|
| 56 |
+
# pipe_14b.to("cpu")
|
| 57 |
|
| 58 |
if randomize_seed:
|
| 59 |
actual_seed = random.randint(0, 1000000)
|