artificialguybr commited on
Commit
43d1e19
1 Parent(s): 8e0d4bf

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +94 -0
utils.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gc
2
+ import os
3
+ import random
4
+ import numpy as np
5
+ import json
6
+ import torch
7
+ from PIL import Image, PngImagePlugin
8
+ from datetime import datetime
9
+ from dataclasses import dataclass
10
+ from typing import Callable, Dict, Optional, Tuple
11
+ from diffusers import (
12
+ DDIMScheduler,
13
+ DPMSolverMultistepScheduler,
14
+ DPMSolverSinglestepScheduler,
15
+ EulerAncestralDiscreteScheduler,
16
+ EulerDiscreteScheduler,
17
+ )
18
+
19
+ MAX_SEED = np.iinfo(np.int32).max
20
+
21
+ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
22
+ if randomize_seed:
23
+ seed = random.randint(0, MAX_SEED)
24
+ return seed
25
+
26
+ def seed_everything(seed: int) -> torch.Generator:
27
+ torch.manual_seed(seed)
28
+ torch.cuda.manual_seed_all(seed)
29
+ np.random.seed(seed)
30
+ generator = torch.Generator()
31
+ generator.manual_seed(seed)
32
+ return generator
33
+
34
+ def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]:
35
+ if aspect_ratio == "Custom":
36
+ return None
37
+ width, height = aspect_ratio.split(" x ")
38
+ return int(width), int(height)
39
+
40
+ def aspect_ratio_handler(aspect_ratio: str, custom_width: int, custom_height: int) -> Tuple[int, int]:
41
+ if aspect_ratio == "Custom":
42
+ return custom_width, custom_height
43
+ else:
44
+ width, height = parse_aspect_ratio(aspect_ratio)
45
+ return width, height
46
+
47
+ def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
48
+ scheduler_factory_map = {
49
+ "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
50
+ "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
51
+ "DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"),
52
+ "Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
53
+ "Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(scheduler_config),
54
+ "DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
55
+ }
56
+ return scheduler_factory_map.get(name, lambda: None)()
57
+
58
+ def free_memory() -> None:
59
+ torch.cuda.empty_cache()
60
+ gc.collect()
61
+
62
+ def common_upscale(samples: torch.Tensor, width: int, height: int, upscale_method: str) -> torch.Tensor:
63
+ return torch.nn.functional.interpolate(samples, size=(height, width), mode=upscale_method)
64
+
65
+ def upscale(samples: torch.Tensor, upscale_method: str, scale_by: float) -> torch.Tensor:
66
+ width = round(samples.shape[3] * scale_by)
67
+ height = round(samples.shape[2] * scale_by)
68
+ return common_upscale(samples, width, height, upscale_method)
69
+
70
+ def preprocess_image_dimensions(width, height):
71
+ if width % 8 != 0:
72
+ width = width - (width % 8)
73
+ if height % 8 != 0:
74
+ height = height - (height % 8)
75
+ return width, height
76
+
77
+ def save_image(image, metadata, output_dir):
78
+ current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
79
+ os.makedirs(output_dir, exist_ok=True)
80
+ filename = f"image_{current_time}.png"
81
+ filepath = os.path.join(output_dir, filename)
82
+
83
+ metadata_str = json.dumps(metadata)
84
+ info = PngImagePlugin.PngInfo()
85
+ info.add_text("metadata", metadata_str)
86
+ image.save(filepath, "PNG", pnginfo=info)
87
+ return filepath
88
+
89
+ def is_google_colab():
90
+ try:
91
+ import google.colab
92
+ return True
93
+ except:
94
+ return False