tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

add stable diff utils

Browse Source
This commit is contained in:
terrance 2023-02-17 00:16:25 +00:00
parent 52ab319953
commit a474aa3fb8
1 changed files with 61 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

61
nerfstudio/generative/stable_diffusion_utils.py Normal file
View File

@ -0,0 +1,61 @@
# Copyright 2022 The Nerfstudio Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Stable diffusion utils
"""
def get_text_embedding(batch, prompting_type, text_embeddings, sd_device):
""""""
base = text_embeddings["base_text_embedding"]
front = text_embeddings["front_text_embedding"]
side = text_embeddings["side_text_embedding"]
top = text_embeddings["top_text_embedding"]
back = text_embeddings["back_text_embedding"]
assert prompting_type in ("location_based", "interpolated", "base")
if prompting_type == "location_based":
if batch["vertical"] < 40:
text_embedding = top
elif batch["central"] > 315 or batch["central"] <= 45:
text_embedding = front
elif batch["central"] > 45 and batch["central"] <= 135:
text_embedding = side
elif batch["central"] > 135 and batch["central"] <= 225:
text_embedding = back
else: # batch["central"] > 225 and batch["central"] <= 315:
text_embedding = side
elif prompting_type == "interpolated":
horiz = batch["central"].to(sd_device)
vert = max(batch["vertical"].to(sd_device), 0)
if batch["central"] > 0 or batch["central"] <= 90:
text_embedding = (horiz) * side + (90 - horiz) * front
elif batch["central"] > 90 and batch["central"] <= 180:
text_embedding = (horiz - 90) * back + (180 - horiz) * side
elif batch["central"] > 180 and batch["central"] <= 270:
text_embedding = (horiz - 180) * side + (270 - horiz) * back
else: # batch["central"] > 270 and batch["central"] <= 360:
text_embedding = (horiz - 270) * front + (360 - horiz) * side
text_embedding = text_embedding / 90.0
text_embedding = (vert * text_embedding + (90 - vert) * top) / 90.0
elif prompting_type == "base":
text_embedding = base
return text_embedding