Tasha Upchurch
447e9403ae
This commit introduces a series of refactoring and enhancements across multiple modules, improving code readability, consistency, and adding new functionalities. Highlights: 1. `tests/test_hf_diffusers_nodes.py`: - Removed unnecessary whitespace. - Improved variable assignment consistency (e.g., `ret=ret[0]` to `ret = ret[0]`). - Added a new test case `test_create_unet2d_model` for Unet2d model creation. 2. `server_endpoints.py`: - Cleaned up commented-out code and added more descriptive comments. - Refactored the construction of `real_git_path` for better readability and future maintenance. - Implemented minor formatting improvements for better code structure. 3. `basic.py`: - Introduced a new class `TupleNode` to handle tuple data types. - Modified `StrBreakout` to include dynamic typing for list conversion. - Enhanced code readability with better variable naming and added comments for clarity. 4. `pytorch_nodes.py`: - Added import statements within functions for `torch`. - New classes `SaveTorchTensor` and `LoadTorchTensor` added for saving and loading torch tensors. - `GridSearchTraining` class updated with minor formatting changes. 5. `image.py`: - In `FloodFillNode`, removed a redundant category assignment. - Added two new classes: `DecorateHeightMap` and `RandomImageFromFolder`. These classes provide functionalities for decorating height maps with images and selecting random images from a folder, respectively. 6. `hf_diffusers_nodes.py`: - Fixed formatting in tensor format options and added a new line for clarity. - New class `CreateUnet2dModel` added for creating a UNet2D model with customizable parameters. - Introduced `TrainUnet2dModel` for training a UNet2D model with a dataset, including detailed setup for the training process. 7. `config.yaml`: - Added a new base path for `gitio` to centralize the configuration. These changes aim to improve the maintainability of the codebase, streamline testing processes, and expand the functionalities within the modules, particularly related to image processing and model training in PyTorch and diffusers.
85 lines
2.3 KiB
Python
85 lines
2.3 KiB
Python
import pytest
|
|
from ..hf_diffusers_nodes import DDPMPipline, SampleDDPMPipline, TensorToImage
|
|
from ..hf_diffusers_nodes import LoadHFDataset
|
|
|
|
|
|
# when testing set "UNIT_TEST" to true in the environment to avoid loading all of ComfyUI
|
|
def test_DDPMPipline():
|
|
p_node = DDPMPipline()
|
|
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
|
|
"cuda")
|
|
assert ret is not None
|
|
|
|
|
|
def test_SampleDDPMPipline():
|
|
p_node = DDPMPipline()
|
|
s_node = SampleDDPMPipline()
|
|
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
|
|
"cuda")
|
|
ret = ret[0]
|
|
assert ret is not None
|
|
ret = s_node.sample(ret,
|
|
1,
|
|
25
|
|
)
|
|
assert ret is not None
|
|
|
|
|
|
def test_TensorToImage():
|
|
p_node = DDPMPipline()
|
|
s_node = SampleDDPMPipline()
|
|
t_node = TensorToImage()
|
|
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
|
|
"cuda")
|
|
ret = ret[0]
|
|
assert ret is not None
|
|
ret = s_node.sample(ret,
|
|
1,
|
|
25
|
|
)
|
|
ret = ret[0]
|
|
assert ret is not None
|
|
ret = t_node.tensor_to_image(ret, "HWC")
|
|
ret = ret[0]
|
|
assert ret is not None
|
|
# now convert to pillow image
|
|
from PIL import Image
|
|
import numpy as np
|
|
# ret is float32 0-1 torch tensor we need to convert it to numpy and then to a pillow image
|
|
|
|
ret = ret.squeeze(0)
|
|
ret = ret.squeeze(0)
|
|
ret = ret.numpy()
|
|
ret = ret * 255
|
|
ret = np.uint8(ret)
|
|
ret = Image.fromarray(ret)
|
|
|
|
ret.show()
|
|
|
|
|
|
def test_LoadHFDataset():
|
|
p_node = LoadHFDataset()
|
|
ret = p_node.load_from_hf("huggan/smithsonian_butterflies_subset")
|
|
assert ret is not None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|
|
|
|
|
|
def test_create_unet2d_model():
|
|
from ..hf_diffusers_nodes import CreateUnet2dModel
|
|
p_node = CreateUnet2dModel()
|
|
ret = p_node.create_unet2d_model(3,
|
|
3,
|
|
1,
|
|
1,
|
|
1,
|
|
1,
|
|
)
|
|
assert ret is not None
|
|
|
|
|
|
|
|
|