fix some minor issue

mit-han-lab · Mar 5, 2024 · 042ae44 · 042ae44
1 parent bcd91f2
commit 042ae44
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -79,9 +79,11 @@ from distrifuser.utils import DistriConfig
 
 distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)
 pipeline = DistriSDXLPipeline.from_pretrained(
-    distri_config=distri_config, pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"
+    distri_config=distri_config,
+    pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0",
+    variant="fp16",
+    use_safetensors=True,
 )
-pipeline.prepare()
 
 pipeline.set_progress_bar_config(disable=distri_config.rank != 0)
 image = pipeline(

diff --git a/distrifuser/__version__.py b/distrifuser/__version__.py
@@ -1 +1 @@
-__version__ = "0.0.0beta0"
+__version__ = "0.0.0beta1"
diff --git a/distrifuser/pipelines.py b/distrifuser/pipelines.py
@@ -14,6 +14,8 @@ def __init__(self, pipeline: StableDiffusionXLPipeline, module_config: DistriCon
 
         self.static_inputs = None
 
+        self.prepare()
+
     @staticmethod
     def from_pretrained(distri_config: DistriConfig, **kwargs):
         device = distri_config.device

diff --git a/scripts/generate_coco.py b/scripts/generate_coco.py
@@ -91,7 +91,6 @@ def main():
         use_safetensors=True,
         scheduler=scheduler,
     )
-    pipeline.prepare()
     pipeline.set_progress_bar_config(disable=distri_config.rank != 0, position=1, leave=False)
 
     if args.output_root is None:

diff --git a/scripts/run_sdxl.py b/scripts/run_sdxl.py
@@ -109,7 +109,6 @@ def main():
         use_safetensors=True,
         scheduler=scheduler,
     )
-    pipeline.prepare()
 
     if args.mode == "generation":
         assert args.output_path is not None

diff --git a/scripts/sdxl_example.py b/scripts/sdxl_example.py
@@ -5,9 +5,11 @@
 
 distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)
 pipeline = DistriSDXLPipeline.from_pretrained(
-    distri_config=distri_config, pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"
+    distri_config=distri_config,
+    pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0",
+    variant="fp16",
+    use_safetensors=True,
 )
-pipeline.prepare()
 
 pipeline.set_progress_bar_config(disable=distri_config.rank != 0)
 image = pipeline(
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		__version__ = "0.0.0beta0"
		__version__ = "0.0.0beta1"