adding base t2i-pipeline for generating 512*1024 panorama image from text input
This commit is contained in:
mengyang.fmy
Generated
+3
@@ -0,0 +1,3 @@
|
||||
# Default ignored files
|
||||
/shelf/
|
||||
/workspace.xml
|
||||
Generated
+8
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<module type="PYTHON_MODULE" version="4">
|
||||
<component name="NewModuleRootManager">
|
||||
<content url="file://$MODULE_DIR$" />
|
||||
<orderEntry type="jdk" jdkName="Python 3.10 (ComfyUI) (2)" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
</component>
|
||||
</module>
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
<component name="InspectionProjectProfileManager">
|
||||
<profile version="1.0">
|
||||
<option name="myName" value="Project Default" />
|
||||
<inspection_tool class="PyInterpreterInspection" enabled="false" level="WARNING" enabled_by_default="false" />
|
||||
<inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
|
||||
<option name="ignoredPackages">
|
||||
<value>
|
||||
<list size="3">
|
||||
<item index="0" class="java.lang.String" itemvalue="hydra-core" />
|
||||
<item index="1" class="java.lang.String" itemvalue="omegaconf" />
|
||||
<item index="2" class="java.lang.String" itemvalue="sentencepiece" />
|
||||
</list>
|
||||
</value>
|
||||
</option>
|
||||
</inspection_tool>
|
||||
<inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
|
||||
<option name="ignoredErrors">
|
||||
<list>
|
||||
<option value="N806" />
|
||||
<option value="N802" />
|
||||
</list>
|
||||
</option>
|
||||
</inspection_tool>
|
||||
</profile>
|
||||
</component>
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
<component name="InspectionProjectProfileManager">
|
||||
<settings>
|
||||
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||
<version value="1.0" />
|
||||
</settings>
|
||||
</component>
|
||||
Generated
+10
@@ -0,0 +1,10 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="MarkdownEnhProjectSettings">
|
||||
<AnnotatorSettings targetHasSpaces="true" linkCaseMismatch="true" wikiCaseMismatch="true" wikiLinkHasDashes="true" notUnderWikiHome="true" targetNotWikiPageExt="true" notUnderSourceWikiHome="true" targetNameHasAnchor="true" targetPathHasAnchor="true" wikiLinkHasSlash="true" wikiLinkHasSubdir="true" wikiLinkHasOnlyAnchor="true" linkTargetsWikiHasExt="true" linkTargetsWikiHasBadExt="true" notUnderSameRepo="true" targetNotUnderVcs="false" linkNeedsExt="true" linkHasBadExt="true" linkTargetNeedsExt="true" linkTargetHasBadExt="true" wikiLinkNotInWiki="true" imageTargetNotInRaw="true" repoRelativeAcrossVcsRoots="true" multipleWikiTargetsMatch="true" unresolvedLinkReference="true" linkIsIgnored="true" anchorIsIgnored="true" anchorIsUnresolved="true" anchorLineReferenceIsUnresolved="true" anchorLineReferenceFormat="true" anchorHasDuplicates="true" abbreviationDuplicates="true" abbreviationNotUsed="true" attributeIdDuplicateDefinition="true" attributeIdNotUsed="true" footnoteDuplicateDefinition="true" footnoteUnresolved="true" footnoteDuplicates="true" footnoteNotUsed="true" macroDuplicateDefinition="true" macroUnresolved="true" macroDuplicates="true" macroNotUsed="true" referenceDuplicateDefinition="true" referenceUnresolved="true" referenceDuplicates="true" referenceNotUsed="true" referenceUnresolvedNumericId="true" enumRefDuplicateDefinition="true" enumRefUnresolved="true" enumRefDuplicates="true" enumRefNotUsed="true" enumRefLinkUnresolved="true" enumRefLinkDuplicates="true" simTocUpdateNeeded="true" simTocTitleSpaceNeeded="true" />
|
||||
<HtmlExportSettings updateOnSave="false" parentDir="" targetDir="" cssDir="css" scriptDir="js" plainHtml="false" imageDir="" copyLinkedImages="false" imagePathType="0" targetPathType="2" targetExt="" useTargetExt="false" noCssNoScripts="false" useElementStyleAttribute="false" linkToExportedHtml="true" exportOnSettingsChange="true" regenerateOnProjectOpen="false" linkFormatType="HTTP_ABSOLUTE" />
|
||||
<LinkMapSettings>
|
||||
<textMaps />
|
||||
</LinkMapSettings>
|
||||
</component>
|
||||
</project>
|
||||
Generated
+55
@@ -0,0 +1,55 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="FlexmarkProjectSettings">
|
||||
<FlexmarkHtmlSettings flexmarkSpecExampleRendering="0" flexmarkSpecExampleRenderHtml="false">
|
||||
<flexmarkSectionLanguages>
|
||||
<option name="1" value="Markdown" />
|
||||
<option name="2" value="HTML" />
|
||||
<option name="3" value="flexmark-ast:1" />
|
||||
</flexmarkSectionLanguages>
|
||||
</FlexmarkHtmlSettings>
|
||||
</component>
|
||||
<component name="MarkdownProjectSettings">
|
||||
<PreviewSettings splitEditorLayout="SPLIT" splitEditorPreview="PREVIEW" useGrayscaleRendering="false" zoomFactor="1.0" maxImageWidth="0" synchronizePreviewPosition="true" highlightPreviewType="LINE" highlightFadeOut="5" highlightOnTyping="true" synchronizeSourcePosition="true" verticallyAlignSourceAndPreviewSyncPosition="true" showSearchHighlightsInPreview="true" showSelectionInPreview="true" lastLayoutSetsDefault="false">
|
||||
<PanelProvider>
|
||||
<provider providerId="com.vladsch.md.nav.editor.swing.html.panel" providerName="Default - Swing" />
|
||||
</PanelProvider>
|
||||
</PreviewSettings>
|
||||
<ParserSettings gitHubSyntaxChange="false" correctedInvalidSettings="false" emojiShortcuts="1" emojiImages="0">
|
||||
<PegdownExtensions>
|
||||
<option name="ATXHEADERSPACE" value="true" />
|
||||
<option name="FENCED_CODE_BLOCKS" value="true" />
|
||||
<option name="INTELLIJ_DUMMY_IDENTIFIER" value="true" />
|
||||
<option name="RELAXEDHRULES" value="true" />
|
||||
<option name="STRIKETHROUGH" value="true" />
|
||||
<option name="TABLES" value="true" />
|
||||
<option name="TASKLISTITEMS" value="true" />
|
||||
</PegdownExtensions>
|
||||
<ParserOptions>
|
||||
<option name="COMMONMARK_LISTS" value="true" />
|
||||
<option name="EMOJI_SHORTCUTS" value="true" />
|
||||
<option name="GFM_TABLE_RENDERING" value="true" />
|
||||
<option name="PRODUCTION_SPEC_PARSER" value="true" />
|
||||
<option name="SIM_TOC_BLANK_LINE_SPACER" value="true" />
|
||||
</ParserOptions>
|
||||
</ParserSettings>
|
||||
<HtmlSettings headerTopEnabled="false" headerBottomEnabled="false" bodyTopEnabled="false" bodyBottomEnabled="false" addPageHeader="false" addAnchorLinks="true" anchorLinksWrapText="false" imageUriSerials="false" addDocTypeHtml="true" noParaTags="false" defaultUrlTitle="false" migratedPlantUml="false" migratedAnchorLinks="false" plantUmlConversion="0">
|
||||
<GeneratorProvider>
|
||||
<provider providerId="com.vladsch.md.nav.editor.text.html.generator" providerName="Unmodified HTML Generator" />
|
||||
</GeneratorProvider>
|
||||
<headerTop />
|
||||
<headerBottom />
|
||||
<bodyTop />
|
||||
<bodyBottom />
|
||||
<fencedCodeConversions />
|
||||
</HtmlSettings>
|
||||
<CssSettings previewScheme="UI_SCHEME" cssUri="" isCssUriEnabled="false" isCssUriSerial="true" isCssTextEnabled="false" isDynamicPageWidth="true">
|
||||
<StylesheetProvider>
|
||||
<provider providerId="com.vladsch.md.nav.editor.text.html.css" providerName="No Stylesheet" />
|
||||
</StylesheetProvider>
|
||||
<ScriptProviders />
|
||||
<cssText />
|
||||
<cssUriHistory />
|
||||
</CssSettings>
|
||||
</component>
|
||||
</project>
|
||||
Generated
+4
@@ -0,0 +1,4 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (ComfyUI) (2)" project-jdk-type="Python SDK" />
|
||||
</project>
|
||||
Generated
+8
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/.idea/Diffusion360_ComfyUI.iml" filepath="$PROJECT_DIR$/.idea/Diffusion360_ComfyUI.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
Generated
+6
@@ -0,0 +1,6 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
@@ -0,0 +1,69 @@
|
||||
import comfy
|
||||
from comfy.ldm.models.autoencoder import AutoencoderKL
|
||||
from comfy.samplers import KSAMPLER
|
||||
from .utils import decode_tiled_blended, decode_tiled_blended_, common_ksampler
|
||||
|
||||
|
||||
class VAEDecodeTiledBlended:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"samples": ("LATENT", ), "vae": ("VAE", ),
|
||||
"tile_size": ("INT", {"default": 512, "min": 320, "max": 4096, "step": 64})
|
||||
}}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "decode"
|
||||
|
||||
CATEGORY = "Diffusion360"
|
||||
|
||||
def decode(self, vae, samples, tile_size):
|
||||
vae.decode_tiled = decode_tiled_blended.__get__(vae, AutoencoderKL)
|
||||
vae.decode_tiled_blended_ = decode_tiled_blended_.__get__(vae, AutoencoderKL)
|
||||
return (vae.decode_tiled(samples["samples"], tile_x=tile_size // 8, tile_y=tile_size // 8, ), )
|
||||
|
||||
|
||||
class Diffusion360Sampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"model": ("MODEL",),
|
||||
"add_noise": (["enable", "disable"], ),
|
||||
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
||||
"sampler_name": (['euler_blend'], ),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"latent_image": ("LATENT", ),
|
||||
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
||||
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
||||
"return_with_leftover_noise": (["disable", "enable"], ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "sample"
|
||||
|
||||
CATEGORY = "Diffusion360"
|
||||
|
||||
def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
|
||||
force_full_denoise = True
|
||||
if return_with_leftover_noise == "enable":
|
||||
force_full_denoise = False
|
||||
disable_noise = False
|
||||
if add_noise == "disable":
|
||||
disable_noise = True
|
||||
return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"VAEDecodeTiledBlended": VAEDecodeTiledBlended,
|
||||
"Diffusion360Sampler": Diffusion360Sampler,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VAEDecodeTiledBlended": "VAE Decode (Tiled Blended)",
|
||||
"Diffusion360Sampler": "Diffusion360Sampler",
|
||||
}
|
||||
@@ -1,2 +1,19 @@
|
||||
# Diffusion360_ComfyUI
|
||||
ComfyUI plugin of https://github.com/ArcherFMY/SD-T2I-360PanoImage
|
||||
|
||||
## features
|
||||
- [x] base t2i-pipeline for generating 512*1024 panorama image from text input
|
||||
- [ ] SR-pipeline for higher-resolution panorama image generation
|
||||
|
||||
## Installation
|
||||
1. Install ComfyUI following https://github.com/comfyanonymous/ComfyUI
|
||||
2. Install this plugin by the following commands.
|
||||
```
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/ArcherFMY/Diffusion360_ComfyUI
|
||||
```
|
||||
|
||||
## Usage
|
||||
add 'Diffusion360Sampler' and 'VAE Decode (Tiled Blended)' from 'Diffusion360' category, and the pipeline looks like
|
||||
|
||||
|
||||
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
from .Diffusion360_nodes import VAEDecodeTiledBlended, Diffusion360Sampler
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"VAEDecodeTiledBlended": VAEDecodeTiledBlended,
|
||||
"Diffusion360Sampler": Diffusion360Sampler,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VAEDecodeTiledBlended": "VAE Decode (Tiled Blended)",
|
||||
"Diffusion360Sampler": "Diffusion360Sampler",
|
||||
}
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 477 KiB |
@@ -0,0 +1,194 @@
|
||||
import torch
|
||||
import comfy
|
||||
from comfy import model_management
|
||||
from tqdm.auto import trange
|
||||
import comfy.k_diffusion.utils as utils
|
||||
import latent_preview
|
||||
from comfy.samplers import KSAMPLER, ksampler, CFGGuider
|
||||
from comfy.extra_samplers import uni_pc
|
||||
|
||||
|
||||
def sampler_object(name):
|
||||
if name == "uni_pc":
|
||||
sampler = KSAMPLER(uni_pc.sample_unipc)
|
||||
elif name == "uni_pc_bh2":
|
||||
sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
|
||||
elif name == "ddim":
|
||||
sampler = ksampler("euler", inpaint_options={"random": True})
|
||||
else:
|
||||
sampler = ksampler(name)
|
||||
return sampler
|
||||
|
||||
|
||||
def sample_(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||
cfg_guider = CFGGuider(model)
|
||||
cfg_guider.set_conds(positive, negative)
|
||||
cfg_guider.set_cfg(cfg)
|
||||
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
|
||||
|
||||
|
||||
def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None,
|
||||
force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
|
||||
if sigmas is None:
|
||||
sigmas = self.sigmas
|
||||
|
||||
if last_step is not None and last_step < (len(sigmas) - 1):
|
||||
sigmas = sigmas[:last_step + 1]
|
||||
if force_full_denoise:
|
||||
sigmas[-1] = 0
|
||||
|
||||
if start_step is not None:
|
||||
if start_step < (len(sigmas) - 1):
|
||||
sigmas = sigmas[start_step:]
|
||||
else:
|
||||
if latent_image is not None:
|
||||
return latent_image
|
||||
else:
|
||||
return torch.zeros_like(noise)
|
||||
|
||||
sampler = sampler_object(self.sampler)
|
||||
sampler.sampler_function = sample_euler_blend
|
||||
|
||||
return sample_(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options,
|
||||
latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar,
|
||||
seed=seed)
|
||||
|
||||
|
||||
def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
|
||||
latent_image = latent["samples"]
|
||||
if disable_noise:
|
||||
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
|
||||
else:
|
||||
batch_inds = latent["batch_index"] if "batch_index" in latent else None
|
||||
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
|
||||
|
||||
noise_mask = None
|
||||
if "noise_mask" in latent:
|
||||
noise_mask = latent["noise_mask"]
|
||||
|
||||
callback = latent_preview.prepare_callback(model, steps)
|
||||
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
||||
|
||||
sampler = comfy.samplers.KSampler(model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
|
||||
sampler.sample = sample.__get__(sampler, KSAMPLER)
|
||||
samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=None, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||
samples = samples.to(comfy.model_management.intermediate_device())
|
||||
|
||||
out = latent.copy()
|
||||
out["samples"] = samples
|
||||
return (out, )
|
||||
|
||||
|
||||
def to_d(x, sigma, denoised):
|
||||
"""Converts a denoiser output to a Karras ODE derivative."""
|
||||
return (x - denoised) / utils.append_dims(sigma, x.ndim)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def blend_h(a, b, blend_extent):
|
||||
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
|
||||
for x in range(blend_extent):
|
||||
b[:, :, :, x] = a[:, :, :, -blend_extent
|
||||
+ x] * (1 - x / blend_extent) + b[:, :, :, x] * (
|
||||
x / blend_extent)
|
||||
return b
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_euler_blend(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
|
||||
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
w = x.shape[-1]
|
||||
for i in trange(len(sigmas) - 1, disable=disable):
|
||||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||||
sigma_hat = sigmas[i] * (gamma + 1)
|
||||
if gamma > 0:
|
||||
eps = torch.randn_like(x) * s_noise
|
||||
x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
|
||||
denoised = model(x, sigma_hat * s_in, **extra_args)
|
||||
d = to_d(x, sigma_hat, denoised)
|
||||
if callback is not None:
|
||||
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
|
||||
dt = sigmas[i + 1] - sigma_hat
|
||||
# Euler method
|
||||
x = x + d * dt
|
||||
x = blend_h(x, x, 4)
|
||||
x = blend_h(x, x, 4)
|
||||
x = x[:, :, :, :w]
|
||||
return x
|
||||
|
||||
|
||||
def decode_tiled_blended_(self, samples, tile_x=64, tile_y=64, overlap=16):
|
||||
steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y,
|
||||
overlap)
|
||||
steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2,
|
||||
tile_y * 2, overlap)
|
||||
steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2,
|
||||
tile_y // 2, overlap)
|
||||
pbar = comfy.utils.ProgressBar(steps)
|
||||
|
||||
decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
|
||||
output = self.process_output(
|
||||
(tiled_scale_blended(samples, decode_fn, tile_x // 2, tile_y * 2, overlap,
|
||||
upscale_amount=self.upscale_ratio, output_device=self.output_device,
|
||||
pbar=pbar) +
|
||||
tiled_scale_blended(samples, decode_fn, tile_x * 2, tile_y // 2, overlap,
|
||||
upscale_amount=self.upscale_ratio, output_device=self.output_device,
|
||||
pbar=pbar) +
|
||||
tiled_scale_blended(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount=self.upscale_ratio,
|
||||
output_device=self.output_device, pbar=pbar))
|
||||
/ 3.0)
|
||||
return output
|
||||
|
||||
|
||||
def decode_tiled_blended(self, samples, tile_x=64, tile_y=64, overlap = 16):
|
||||
model_management.load_model_gpu(self.patcher)
|
||||
output = self.decode_tiled_blended_(samples, tile_x, tile_y, overlap)
|
||||
return output.movedim(1, -1)
|
||||
|
||||
|
||||
def blend_h(a, b, blend_extent):
|
||||
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
|
||||
for x in range(blend_extent):
|
||||
b[:, :, :, x] = a[:, :, :, -blend_extent
|
||||
+ x] * (1 - x / blend_extent) + b[:, :, :, x] * (
|
||||
x / blend_extent)
|
||||
return b
|
||||
|
||||
|
||||
def tiled_scale_blended(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
|
||||
output = torch.empty((samples.shape[0], out_channels, round(samples.shape[2] * upscale_amount), round(samples.shape[3] * upscale_amount)), device=output_device)
|
||||
for b in range(samples.shape[0]):
|
||||
s = samples[b:b+1]
|
||||
out = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
|
||||
out_div = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
|
||||
w = samples.shape[3]
|
||||
samples = torch.cat([samples, samples[:, :, :, :w // 4]], dim=-1)
|
||||
s = samples[b:b + 1]
|
||||
for y in range(0, s.shape[2], tile_y - overlap):
|
||||
# for x in range(0, s.shape[3], tile_x - overlap):
|
||||
# x = max(0, min(s.shape[-1] - overlap, x))
|
||||
y = max(0, min(s.shape[-2] - overlap, y))
|
||||
s_in = s[:,:,y:y+tile_y,:]
|
||||
|
||||
ps = function(s_in).to(output_device)
|
||||
ps = blend_h(
|
||||
ps[:,:,:,w*upscale_amount:],
|
||||
ps[:,:,:,:w*upscale_amount],
|
||||
w // 4 * upscale_amount
|
||||
)
|
||||
mask = torch.ones_like(ps)
|
||||
feather = round(overlap * upscale_amount)
|
||||
for t in range(feather):
|
||||
mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))
|
||||
mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
|
||||
mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
|
||||
mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
|
||||
out[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),:] += ps * mask
|
||||
out_div[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),:] += mask
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
output[b:b+1] = out/out_div
|
||||
return output
|
||||
Reference in New Issue
Block a user