Hugging Face's logo

Spaces:
tungedng2710
/
TonAI-Creative
Running

App Files Community
TonAI-Creative
File size: 4,397 Bytes 
8eca2ce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b4bf53d
 
 
 
8eca2ce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
 
import torch
import GPUtil
import math
import random
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, \
    FluxPipeline, StableDiffusion3Pipeline

MAX_SEED = 2**32 - 1
TEXT_TO_IMAGE_DICTIONARY = {
    # "FLUX.1 [schnell-dev-merged]": {
    #     "path": "sayakpaul/FLUX.1-merged",
    #     "pipeline": FluxPipeline,
    #     "device_map": "balanced"
    # },
    # "FLUX.1 [schnell]": {
    #     "path": "black-forest-labs/FLUX.1-schnell",
    #     "pipeline": FluxPipeline,
    #     "device_map": "balanced"
    # },
    # "FLUX.1 [dev]": {
    #     "path": "black-forest-labs/FLUX.1-dev",
    #     "pipeline": FluxPipeline,
    #     "device_map": "balanced"
    # },
    # "Stable Diffusion 2.1": {
    #     "path": "stabilityai/stable-diffusion-2-1",
    #     "pipeline": StableDiffusionPipeline
    # },
    "Stable Diffusion 3.5 Medium": {
        "backend": "comfyui",
        "path": "stuffs/comfyui_workflow_api/sd3_5_workflow_api.json",
        "device_map": "balanced"
    },
    "Stable Diffusion 3.5 Large": {
        "backend": "comfyui",
        "path": "stuffs/comfyui_workflow_api/sd3_5_workflow_api.json",
        "device_map": "balanced"
    },
    # "Stable Diffusion 3 Medium": {
    #     "path": "stabilityai/stable-diffusion-3-medium-diffusers",
    #     "pipeline": StableDiffusion3Pipeline,
    #     "device_map": "balanced"
    # },
    # "Realistic SDXL": {
    #     "path": "misri/epicrealismXL_v7FinalDestination",
    #     "pipeline": StableDiffusionXLPipeline,
    # },
    # "DreamShaper8 (SD 1.5)": {
    #     "path": "Lykon/dreamshaper-8",
    #     "pipeline": StableDiffusionPipeline,
    # },
    "Anime (SD 1.5)": {
        "path": "../checkpoints/darkSushiMixMix_225D.safetensors",
        "pipeline": StableDiffusionPipeline,
    }
}


def nearest_divisible_by_8(number: int = 1024):
    """
    Auto adjust the number to make it divisible by 8
    """
    lower_multiple = (number // 8) * 8
    upper_multiple = lower_multiple + 8
    if (number - lower_multiple) < (upper_multiple - number):
        return int(lower_multiple)
    else:
        return int(upper_multiple)


def get_gpu_info(width: int = 1024, 
                 height: int = 1024,
                 num_images: int = 1) -> tuple[list, dict]:
    """
    Get available GPUs info

    Parameters:
        - width : generated image's width
        - height : generated image's height
        - num_images : number of generated images per prompt
    Returns:
        - gpu_info and current_max_memory
    """
    gpus = GPUtil.getGPUs()
    gpu_info = []
    current_max_memory = {}
    using_fast_flux = width <= 1280 \
              and height <= 1280 \
              and num_images==1
    for gpu in gpus:
        info = {
            'id': gpu.id,
            'name': gpu.name,
            'driver_version': gpu.driver,
            'total_memory': gpu.memoryTotal,  # In MB
            'available_memory': gpu.memoryFree,  # In MB
            'used_memory': gpu.memoryUsed,  # In MB
            'temperature': gpu.temperature  # In Celsius
        }
        gpu_info.append(info)
        if using_fast_flux:
            current_max_memory[gpu.id] = f"{math.ceil(gpu.memoryFree / 1024)}GB"
        else:
            current_max_memory[gpu.id] = f"{int(gpu.memoryFree / 1024)}GB"


    return gpu_info, current_max_memory


def generate_number():
    """
    Random an integer

    Returns:
        - int: an integer
    """
    return random.randint(0, MAX_SEED)


def assign_gpu(required_vram, width, height, num_images):
    """
    Assign GPU device
    
    Parameters:
        - required_memory (int): minimum VRAM
    
    Returns:
        - torch.device
    """
    gpu_info, _ = get_gpu_info(width, height, num_images)
    device = "cpu"
    for gpu in gpu_info:
        if gpu['available_memory'] >= required_vram:
            device = f"cuda:{gpu['id']}"
    if device == "cpu":
        return device
    return torch.device(device)


# def optimize_flux_pipeline(width, height, num_images):
#     """
    

#     Parameters:
#         - width (int): generated image width
#         - height (int): generated image height
#         - num_images (int): number of generated images per prompt
#     """
#     using_fast_flux = width <= 1280 \
#               and height <= 1280 \
#               and num_images==1
#     return using_fast_flux