File size: 4,189 Bytes
2fdce3c |
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 |
import gradio as gr
import json
from .class_configuration_file import ConfigurationFile
from .class_source_model import SourceModel
from .class_folders import Folders
from .class_basic_training import BasicTraining
from .class_advanced_training import AdvancedTraining
from .class_sample_images import SampleImages
from library.dreambooth_folder_creation_gui import (
gradio_dreambooth_folder_creation_tab,
)
from .common_gui import color_aug_changed
class Dreambooth:
def __init__(
self,
headless: bool = False,
):
self.headless = headless
self.dummy_db_true = gr.Label(value=True, visible=False)
self.dummy_db_false = gr.Label(value=False, visible=False)
self.dummy_headless = gr.Label(value=headless, visible=False)
gr.Markdown('Train a custom model using kohya dreambooth python code...')
# Setup Configuration Files Gradio
self.config = ConfigurationFile(headless)
self.source_model = SourceModel(headless=headless)
with gr.Tab('Folders'):
self.folders = Folders(headless=headless)
with gr.Tab('Parameters'):
self.basic_training = BasicTraining(
learning_rate_value='1e-5',
lr_scheduler_value='cosine',
lr_warmup_value='10',
)
self.full_bf16 = gr.Checkbox(
label='Full bf16', value = False
)
with gr.Accordion('Advanced Configuration', open=False):
self.advanced_training = AdvancedTraining(headless=headless)
self.advanced_training.color_aug.change(
color_aug_changed,
inputs=[self.advanced_training.color_aug],
outputs=[self.basic_training.cache_latents],
)
self.sample = SampleImages()
with gr.Tab('Tools'):
gr.Markdown(
'This section provide Dreambooth tools to help setup your dataset...'
)
gradio_dreambooth_folder_creation_tab(
train_data_dir_input=self.folders.train_data_dir,
reg_data_dir_input=self.folders.reg_data_dir,
output_dir_input=self.folders.output_dir,
logging_dir_input=self.folders.logging_dir,
headless=headless,
)
def save_to_json(self, filepath):
def serialize(obj):
if isinstance(obj, gr.inputs.Input):
return obj.get()
if isinstance(obj, (bool, int, float, str)):
return obj
if isinstance(obj, dict):
return {k: serialize(v) for k, v in obj.items()}
if hasattr(obj, "__dict__"):
return serialize(vars(obj))
return str(obj) # Fallback for objects that can't be serialized
try:
with open(filepath, 'w') as outfile:
print(serialize(vars(self)))
json.dump(serialize(vars(self)), outfile)
except Exception as e:
print(f"Error saving to JSON: {str(e)}")
def load_from_json(self, filepath):
def deserialize(key, value):
if hasattr(self, key):
attr = getattr(self, key)
if isinstance(attr, gr.inputs.Input):
attr.set(value)
elif hasattr(attr, "__dict__"):
for k, v in value.items():
deserialize(k, v)
else:
setattr(self, key, value)
else:
print(f"Warning: {key} not found in the object's attributes.")
try:
with open(filepath) as json_file:
data = json.load(json_file)
for key, value in data.items():
deserialize(key, value)
except FileNotFoundError:
print(f"Error: The file {filepath} was not found.")
except json.JSONDecodeError:
print(f"Error: The file {filepath} could not be decoded as JSON.")
except Exception as e:
print(f"Error loading from JSON: {str(e)}") |