diff --git a/nerfstudio/generative/stable_diffusion_utils.py b/nerfstudio/generative/stable_diffusion_utils.py index 9804eb96c..a02648cf4 100644 --- a/nerfstudio/generative/stable_diffusion_utils.py +++ b/nerfstudio/generative/stable_diffusion_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +<<<<<<< HEAD """Utility functions for Stable Diffusion""" def get_text_embedding(batch, prompting_type, prompts, device): @@ -22,10 +23,26 @@ def get_text_embedding(batch, prompting_type, prompts, device): prompts: dictionary of text embeddings """ +======= +""" +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"] + +>>>>>>> a474aa3fb81d07d55beb27685b4e8d6cd5bf9514 assert prompting_type in ("location_based", "interpolated", "base") if prompting_type == "location_based": if batch["vertical"] < 40: +<<<<<<< HEAD text_embedding = prompts["top_text_embedding"] elif batch["central"] > 315 or batch["central"] <= 45: text_embedding = prompts["front_text_embedding"] @@ -54,5 +71,35 @@ def get_text_embedding(batch, prompting_type, prompts, device): else: text_embedding = prompts["base_text_embedding"] +======= + 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 +>>>>>>> a474aa3fb81d07d55beb27685b4e8d6cd5bf9514 return text_embedding