Skip to content

Commit

Permalink
Remove unnecessary token from Stable Diffusion CLI (#393)
Browse files Browse the repository at this point in the history
* Remove unnecessary token from Stable Diffusion CLI

* Update stable_diffusion_cli.py

Ruff lint
  • Loading branch information
ekzhang authored Aug 7, 2023
1 parent f5ea162 commit 60b9291
Showing 1 changed file with 2 additions and 10 deletions.
12 changes: 2 additions & 10 deletions 06_gpu_and_ml/stable_diffusion/stable_diffusion_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,10 @@
from __future__ import annotations

import io
import os
import time
from pathlib import Path

from modal import Image, Secret, Stub, method
from modal import Image, Stub, method

# All Modal programs need a [`Stub`](/docs/reference/modal.Stub) — an object that acts as a recipe for
# the application. Let's give it a friendly name.
Expand All @@ -59,22 +58,18 @@ def download_models():
import diffusers
import torch

hugging_face_token = os.environ["HUGGINGFACE_TOKEN"]

# Download scheduler configuration. Experiment with different schedulers
# to identify one that works best for your use-case.
scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
model_id,
subfolder="scheduler",
use_auth_token=hugging_face_token,
cache_dir=cache_path,
)
scheduler.save_pretrained(cache_path, safe_serialization=True)

# Downloads all other models.
pipe = diffusers.StableDiffusionPipeline.from_pretrained(
model_id,
use_auth_token=hugging_face_token,
revision="fp16",
torch_dtype=torch.float16,
cache_dir=cache_path,
Expand All @@ -98,10 +93,7 @@ def download_models():
find_links="https://download.pytorch.org/whl/torch_stable.html",
)
.pip_install("xformers", pre=True)
.run_function(
download_models,
secrets=[Secret.from_name("huggingface-secret")],
)
.run_function(download_models)
)
stub.image = image

Expand Down

0 comments on commit 60b9291

Please sign in to comment.