Spaces:
Running
on
Zero
Running
on
Zero
Commit
•
da62732
1
Parent(s):
80cf914
Create utils.py
Browse files
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
|