Spaces:
Sleeping
Sleeping
File size: 9,204 Bytes
810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e 56ab875 810485e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 |
import os
from dotenv import load_dotenv
import torch
import numpy as np
import cv2
from PIL import Image
from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
class ArchIntelligent:
def __init__(self):
# Get private variables from enviroment
load_dotenv()
self.hf_token = os.getenv("HF_TOKEN")
self.style_models = os.getenv("STYLE_MODELS")
self.functional_models= os.getenv("FUNCTION_MODELS")
self.enhancement= os.getenv("REALISM_ENHANCE")
self.controlnet_model= os.getenv("CONTROLNET")
self.base_model = os.getenv("BASEMODEL")
self.model_config = {}
# Configure ControlNet model
controlnet = ControlNetModel.from_pretrained(
self.controlnet_model,
torch_dtype= torch.float16,
cache_dir= r"huggingface_cache",
token= self.hf_token,
variant= 'fp16',
)
self.pipeline= StableDiffusionXLControlNetPipeline.from_pretrained(
self.base_model,
controlnet= controlnet,
torch_dtype= torch.float16,
cache_dir= r"huggingface_cache",
token= self.hf_token,
variant= 'fp16',
)
# Enable memory-efficient optimizations
try:
self.pipeline.enable_xformers_memory_efficient_attention()
self.pipeline.enable_vae_slicing()
self.pipeline.enable_sequential_cpu_offload()
print(f"xFormers enabled\nVAE Slicing mode enabled\nSequential CPU Offload enabled!")
except Exception as e:
print(f"Warning: Some optimizations failed: {e}")
def img2canny(self, input_img):
"""
Processing user's condition image into edge map
Parameters
input_img : PIL image
Returns
PIL image
"""
np_image = np.array(input_img)
# Convert the image into a grayscale image then extract edge map
canny = cv2.cvtColor(np_image, cv2.COLOR_RGB2GRAY)
canny = cv2.resize(canny, (1024, 1024))
canny = cv2.Canny(canny, 100, 200)
canny = Image.fromarray(canny)
return canny
def process_config(self, config: dict):
style_dict = {"Modern": "Modernism", "Minimalism": "Minimalism", "Art Deco": "ArtDeco",
"Art Nouveau": "ArtNouveau", "Baroque": "Baroque", "Brutalist": "Brutalist",
"Classical": "Classical", "Neo-Classical": "Neo-Classical", "Cyberpunk": "Cyberpunk",
"Deconstructivism": "Deconstructivism", "Futurism": "Futurism", "Gothic": "Gothic",
"Neo-Futurism": "Neo-Futurism", "Sustainable": "Sustainable", "Victorian": "Victorian"}
functional_dict = {"Residential": "Modern", "Villa": "Modern", "Office": "Office", "Skyscraper": "SkyScraper",
"Hotel": "Hotel", "School Campus": "SchoolCampus", "Farmhouse": "Farmhouse", "Playground": "PlayGround",
"Park": "Park", "Apartment": "Apartment", "Hospital": "Hospital", "Kindergarten": "KinderGarten",
"Church": "Church", "Container": "Container", "Bridge": "Bridge", "Resort": "Resort", "Airport": "Airport",
"Factory": "Factory", "Stadium": "Stadium", "Temple": "Temple", "Tree House": "TreeHouse"}
styles= config['style_names']
functional= config['functional_names']
season = config['season']
landscape= config['landscape']
weather= config['weather']
day= config['time_of_day']
config['posprompt_2'] = f"(((realistic))), (({styles})), (({functional})), ({landscape}), ({season}), ({weather}), ({day}), (high quality),\
(high resolution), 4k render, detail, beautiful, cinematic lighting, hyper-realistic"
config['negprompt_2'] = "((blurry)), details are low, overlapping, (grainy), multiple angles, deformed structures, unnatural, unrealistic, cartoon, \
anime, (painting), drawing, sketch, gibberish text, logo, noise, jpeg artifacts, mutation, (((worst quality))), ((low quality)), (((low resolution))),\
messy, watermark, signature, cut off, low contrast, underexposed, overexposed, draft, disfigured, ugly, tiling, out of frame"
config["LoRA_style"] = style_dict[styles]
config["LoRA_functional"] = functional_dict[functional]
config['adapter_weights'] = [1.0, 1.0, 0.8]
self.model_config = config
def generate(self):
"""
Generate building image using user's input arguments
"""
# Get user's prompts from dictionary
first_prompt = self.model_config["posprompt_1"]
second_prompt = self.model_config["posprompt_2"]
first_negprompt = self.model_config["negprompt_1"]
second_negprompt = self.model_config["negprompt_2"]
# Get user's image
input_image = self.model_config['image']
# Get ControlNet conditioning scale value
controlnet_condition = self.model_config["condition_scale"]
# Get guidance scale value
guidance_scale = self.model_config["guidance"]
# Get render speed
render_speed = self.model_config["render_speed"]
# Get LoRA weight's name and their corresponding adapter weights
LoRA_style_names = self.model_config['LoRA_style']
LoRA_functional_names = self.model_config['LoRA_functional']
LoRA_enhancement_names = 'Realism'
adapter_weights = self.model_config['adapter_weights']
LoRA_names = [LoRA_style_names, LoRA_functional_names, LoRA_enhancement_names]
self.pipeline.unload_lora_weights()
print(f"\n\nUNLOADED LORA WEIGHTS\n\n")
os.environ['HF_HOME'] = r"huggingface_cache"
self.pipeline.load_lora_weights(
self.style_models,
weight_name= f"{LoRA_style_names}.safetensors",
adapter_name= LoRA_style_names
)
self.pipeline.load_lora_weights(
self.functional_models,
weight_name= f"{LoRA_functional_names}.safetensors",
adapter_name= LoRA_functional_names
)
self.pipeline.load_lora_weights(
self.enhancement,
weight_name= f"realistic.safetensors",
adapter_name= LoRA_enhancement_names
)
print(f"Finished loadded 3 LoRA weights {LoRA_style_names}, {LoRA_functional_names} and {LoRA_enhancement_names}")
self.pipeline.set_adapters(adapter_names= LoRA_names, adapter_weights= adapter_weights)
print(f"Adapted 3 lora weights")
# Transform the image into a depth map that is compatible with ControlNet
conditional_image = self.img2canny(input_image)
# Setup the pipeline then generate image
image = self.pipeline(
prompt= first_prompt,
prompt_2= second_prompt,
negative_prompt= first_negprompt,
negative_prompt_2= second_negprompt,
image= conditional_image,
controlnet_conditioning_scale= controlnet_condition,
num_inference_steps= render_speed,
guidance_scale= guidance_scale
).images[0]
return image
if __name__ == '__main__':
print("Loading")
pipe = ArchIntelligent()
print("Finished") |