Update README.md

README.md

@@ -6,9 +6,21 @@ license: cc-by-nc-4.0
 
 **Demo:**
 ```python
+import random, os
+import numpy as np
 from pathlib import Path
 import torch, torchvision
 
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+
+Path("demo").mkdir(parents=True, exist_ok=True)
+
+    
 from diffusers import StableDiffusionXLPipeline
 pipe = StableDiffusionXLPipeline.from_pretrained("hansyan/perflow-sdxl-base", torch_dtype=torch.float16, use_safetensors=True, variant="v0-fix")
 from src.scheduler_perflow import PeRFlowScheduler
@@ -49,5 +61,4 @@ for cfg_scale in cfg_scale_list:
         cfg_int = int(cfg_scale); cfg_float = int(cfg_scale*10 - cfg_int*10)
         save_name = f'step_{num_inference_steps}_txt{i+1}_cfg{cfg_int}-{cfg_float}.png'
         torchvision.utils.save_image(torchvision.utils.make_grid(samples, nrow = num_img), os.path.join("demo", save_name))
-
 ```