diff --git a/Diffusion360_nodes_diffusers.py b/Diffusion360_nodes_diffusers.py index 7bbbae4..fee8bb5 100644 --- a/Diffusion360_nodes_diffusers.py +++ b/Diffusion360_nodes_diffusers.py @@ -36,7 +36,7 @@ class InputImage: FUNCTION = "load_image" def load_image(self, image): - image_path = os.path.join('custom_nodes/Diffusion360_ComfyUI/data/', image) + image_path = os.path.join('custom_nodes', 'Diffusion360_ComfyUI', 'data', image) img = node_helpers.pillow(Image.open, image_path) output_images = [] @@ -123,7 +123,7 @@ class Diffusion360LoaderText2Pano: CATEGORY = "Diffusion360/diffusers" def load_models(self, model_path): - pipe = Text2360PanoramaImagePipeline(os.path.join('models/diffusers', model_path), torch_dtype=torch.float16) + pipe = Text2360PanoramaImagePipeline(os.path.join('models', 'diffusers', model_path), torch_dtype=torch.float16) return (pipe, ) @@ -184,6 +184,7 @@ class Diffusion360LoaderImage2Pano: CATEGORY = "Diffusion360/diffusers" def load_models(self, model_path): - pipe = Image2360PanoramaImagePipeline(os.path.join('models/diffusers', model_path), torch_dtype=torch.float16) - mask = load_image('custom_nodes/Diffusion360_ComfyUI/data/i2p-mask.jpg') + pipe = Image2360PanoramaImagePipeline(os.path.join('models', 'diffusers', model_path), torch_dtype=torch.float16) + mask_path = os.path.join('custom_nodes', 'Diffusion360_ComfyUI', 'data', 'i2p-mask.jpg') + mask = load_image(mask_path) return (pipe, mask) diff --git a/img2panoimg/image_to_360panorama_image_pipeline.py b/img2panoimg/image_to_360panorama_image_pipeline.py index 1e4be24..5cc8779 100644 --- a/img2panoimg/image_to_360panorama_image_pipeline.py +++ b/img2panoimg/image_to_360panorama_image_pipeline.py @@ -1,6 +1,7 @@ # Copyright © Alibaba, Inc. and its affiliates. import random from typing import Any, Dict +import os import numpy as np import torch @@ -47,10 +48,12 @@ class Image2360PanoramaImagePipeline(DiffusionPipeline): enable_xformers_memory_efficient_attention = kwargs.get( 'enable_xformers_memory_efficient_attention', True) - model_id = model + '/sr-base/' + # model_id = model + '/sr-base/' + model_id = os.path.join(model, 'sr-base') # init i2p model - controlnet = ControlNetModel.from_pretrained(model + '/sd-i2p', torch_dtype=torch.float16) + # controlnet = ControlNetModel.from_pretrained(model + '/sd-i2p', torch_dtype=torch.float16) + controlnet = ControlNetModel.from_pretrained(os.path.join(model, 'sd-i2p'), torch_dtype=torch.float16) self.pipe = StableDiffusionImage2PanoPipeline.from_pretrained( model_id, controlnet=controlnet, torch_dtype=torch_dtype).to(device) @@ -66,8 +69,11 @@ class Image2360PanoramaImagePipeline(DiffusionPipeline): self.pipe.enable_model_cpu_offload() # init controlnet-sr model - base_model_path = model + '/sr-base' - controlnet_path = model + '/sr-control' + # base_model_path = model + '/sr-base' + # controlnet_path = model + '/sr-control' + base_model_path = os.path.join(model, 'sr-base') + controlnet_path = os.path.join(model, 'sr-control') + controlnet = ControlNetModel.from_pretrained( controlnet_path, torch_dtype=torch_dtype) self.pipe_sr = StableDiffusionControlNetImg2ImgPanoPipeline.from_pretrained( @@ -94,7 +100,8 @@ class Image2360PanoramaImagePipeline(DiffusionPipeline): scale=2) netscale = 2 - model_path = model + '/RealESRGAN_x2plus.pth' + # model_path = model + '/RealESRGAN_x2plus.pth' + model_path = os.path.join(model, 'RealESRGAN_x2plus.pth') dni_weight = None self.upsampler = RealESRGANer( diff --git a/node.tar.gz b/node.tar.gz new file mode 100644 index 0000000..d232b07 Binary files /dev/null and b/node.tar.gz differ diff --git a/txt2panoimg/text_to_360panorama_image_pipeline.py b/txt2panoimg/text_to_360panorama_image_pipeline.py index f113d97..924a85e 100644 --- a/txt2panoimg/text_to_360panorama_image_pipeline.py +++ b/txt2panoimg/text_to_360panorama_image_pipeline.py @@ -1,6 +1,7 @@ # Copyright © Alibaba, Inc. and its affiliates. import random from typing import Any, Dict +import os import numpy as np import torch @@ -46,7 +47,8 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline): enable_xformers_memory_efficient_attention = kwargs.get( 'enable_xformers_memory_efficient_attention', True) - model_id = model + '/sd-base/' + # model_id = model + '/sd-base/' + model_id = os.path.join(model, 'sd-base') # init base model self.pipe = StableDiffusionBlendExtendPipeline.from_pretrained( @@ -63,8 +65,11 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline): self.pipe.enable_model_cpu_offload() # init controlnet-sr model - base_model_path = model + '/sr-base' - controlnet_path = model + '/sr-control' + # base_model_path = model + '/sr-base' + # controlnet_path = model + '/sr-control' + base_model_path = os.path.join(model, 'sr-base') + controlnet_path = os.path.join(model, 'sr-control') + controlnet = ControlNetModel.from_pretrained( controlnet_path, torch_dtype=torch_dtype) self.pipe_sr = StableDiffusionControlNetImg2ImgPanoPipeline.from_pretrained( @@ -91,7 +96,8 @@ class Text2360PanoramaImagePipeline(DiffusionPipeline): scale=2) netscale = 2 - model_path = model + '/RealESRGAN_x2plus.pth' + # model_path = model + '/RealESRGAN_x2plus.pth' + model_path = os.path.join(model, 'RealESRGAN_x2plus.pth') dni_weight = None self.upsampler = RealESRGANer(