Spaces:
Running
Running
import gradio as gr | |
import torch | |
from diffusers import StableDiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel, EulerDiscreteScheduler | |
from jinja2 import Template | |
import numpy as np | |
import cv2 | |
from PIL import Image | |
import time | |
# Load models from Hugging Face | |
sd_model_id = "bhoomikagp/sd2-interior-model-version2" ## test | |
# sd_model_id = "bhoomikagp/sd3-interior-model" ## SD3 model issue loading | |
controlnet_model_id = "lllyasviel/sd-controlnet-mlsd" | |
# Load Stable Diffusion pipeline CUDA | |
scheduler = EulerDiscreteScheduler.from_pretrained(sd_model_id, subfolder="scheduler") | |
#sd_pipeline = StableDiffusionPipeline.from_pretrained(sd_model_id, torch_dtype=torch.float16,scheduler=scheduler).to("cuda") | |
# Load ControlNet and Stable Diffusion ControlNet pipeline | |
""" | |
controlnet = ControlNetModel.from_pretrained(controlnet_model_id, torch_dtype=torch.float16).to("cuda") | |
controlnet_pipeline = StableDiffusionControlNetPipeline.from_pretrained( | |
sd_model_id, | |
controlnet=controlnet, | |
scheduler=scheduler, | |
torch_dtype=torch.float16 | |
).to("cuda") | |
""" | |
#sd_model_id = "stabilityai/stable-diffusion-2-1" | |
#scheduler = DPMSolverMultistepScheduler.from_pretrained(sd_model_id, subfolder="scheduler") | |
# Check if CUDA is available | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
print(f"Using device: {device}") | |
# Initialize the pipeline with appropriate device and dtype | |
torch_dtype = torch.float16 if device == "cuda" else torch.float32 | |
sd_pipeline = StableDiffusionPipeline.from_pretrained( | |
sd_model_id, | |
torch_dtype=torch_dtype, | |
scheduler=scheduler | |
).to(device) | |
# choices lists | |
option_choices = ["living_room", "bedroom", "kitchen"] | |
adj_1_choices = [ | |
"spacious", "cozy", "minimalist", "elegant", "modern", "rustic", | |
"luxurious", "inviting", "airy", "sophisticated", "bright", "warm", | |
"serene", "chic", "contemporary" | |
] | |
architecture_style_choices = [ | |
"modern", "contemporary", "traditional", "industrial", "scandinavian", | |
"mid-century", "colonial", "art deco", "neo-classical", "mediterranean", | |
"gothic", "baroque", "japanese", "brutalist", "tropical" | |
] | |
aesthetic_choices = [ | |
"bohemian", "vintage", "minimalist", "luxurious", "eclectic", | |
"mid-century", "art deco", "modern farmhouse", "industrial chic", | |
"shabby chic", "rustic elegance", "coastal", "urban", "transitional", | |
"contemporary classic" | |
] | |
primary_color_choices = [ | |
"blue", "green", "beige", "grey", "white", "black", | |
"cream", "brown", "taupe", "burgundy", "mustard", | |
"terracotta", "olive", "peach", "navy" | |
] | |
wood_finish_choices = [ | |
"dark oak", "walnut", "mahogany", "teak", "maple", | |
"cherry", "pine", "birch", "ash", "rosewood", | |
"ebony", "cedar", "hickory", "elm", "red oak" | |
] | |
wall_color_choices = [ | |
"cream", "off-white", "charcoal", "sage green", "navy blue", | |
"taupe", "light grey", "soft pink", "mustard yellow", "deep teal", | |
"warm beige", "pearl white", "slate blue", "coral", "mint green" | |
] | |
tiles_choices = [ | |
"marble", "ceramic", "porcelain", "slate", "wooden-look", | |
"mosaic", "granite", "terracotta", "cement", "quartz", | |
"limestone", "onyx", "travertine", "glass", "encaustic" | |
] | |
# Templates for each room type | |
templates = { | |
"living_room": """ | |
High quality, High resolution, Interior, Architecture, Revit, Autocad, Realistic 3D Render, vray, lumion, raytracing, | |
of a {{ adj_1 }} living room, in {{ architecture_style }} architecture, with {{ aesthetic }} style, | |
painted in {{ wall_color }} with {{ primary_color }} accents, {{ wood_finish }} wood finishes, and {{ tiles }} flooring. | |
""", | |
"bedroom": """ | |
High quality, High resolution, Interior, Architecture, Revit, Autocad, Realistic 3D Render, vray, lumion, raytracing, | |
of a {{ adj_1 }} bedroom, in {{ architecture_style }} architecture, with {{ aesthetic }} style, | |
painted in {{ wall_color }} with {{ primary_color }} accents, {{ wood_finish }} wood finishes, and {{ tiles }} flooring. | |
""", | |
"kitchen": """ | |
High quality, High resolution, Interior, Architecture, Revit, Autocad, Realistic 3D Render, vray, lumion, raytracing, | |
of a {{ adj_1 }} kitchen, in {{ architecture_style }} architecture, with {{ aesthetic }} style, | |
painted in {{ wall_color }} with {{ primary_color }} accents, {{ wood_finish }} cabinetry, and {{ tiles }} flooring. | |
""" | |
} | |
def generate_prompt(option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles): | |
# Validate inputs | |
if adj_1 not in adj_1_choices: | |
raise ValueError(f"Invalid adjective: '{adj_1}'. Accepted options are {adj_1_choices}") | |
if architecture_style not in architecture_style_choices: | |
raise ValueError(f"Invalid architecture style: '{architecture_style}'. Accepted options are {architecture_style_choices}") | |
if aesthetic not in aesthetic_choices: | |
raise ValueError(f"Invalid aesthetic style: '{aesthetic}'. Accepted options are {aesthetic_choices}") | |
if primary_color not in primary_color_choices: | |
raise ValueError(f"Invalid primary color: '{primary_color}'. Accepted options are {primary_color_choices}") | |
if wood_finish not in wood_finish_choices: | |
raise ValueError(f"Invalid wood finish: '{wood_finish}'. Accepted options are {wood_finish_choices}") | |
if wall_color not in wall_color_choices: | |
raise ValueError(f"Invalid wall color: '{wall_color}'. Accepted options are {wall_color_choices}") | |
if tiles not in tiles_choices: | |
raise ValueError(f"Invalid tiles: '{tiles}'. Accepted options are {tiles_choices}") | |
# Select the template based on the room type option | |
template_str = templates.get(option.lower()) | |
if not template_str: | |
raise ValueError(f"Invalid option: '{option}'. Available options are {option_choices}") | |
# Render the template | |
template = Template(template_str) | |
return template.render( | |
adj_1=adj_1, | |
architecture_style=architecture_style, | |
aesthetic=aesthetic, | |
primary_color=primary_color, | |
wood_finish=wood_finish, | |
wall_color=wall_color, | |
tiles=tiles | |
) | |
# Function to generate initial image | |
def generate_initial_image(n_images=1, option='living_room', adj_1="spacious", architecture_style="modern", aesthetic="minimalist", | |
primary_color="neutral", wood_finish="oak", wall_color="off-white", tiles="marble"): | |
# Generate the prompt from choices | |
gen_prompt = generate_prompt(option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles) | |
images = [] # List to store generated images along with seeds | |
for i in range(n_images): | |
# Generate a random seed for reproducibility | |
seed = torch.randint(0, 2**32, (1,)).item() | |
generator = torch.manual_seed(seed) | |
# Generate the image with the prompt and Stable Diffusion pipeline | |
cfg = 8 # Configuring guidance scale | |
steps = 30 # Number of inference steps | |
width, height = 640, 512 | |
# pipeline without controlnet | |
img = sd_pipeline(prompt=gen_prompt, guidance_scale=cfg, num_inference_steps=steps, width=width, height=height, generator=generator).images[0] | |
# Append the seed and image to the list for tracking | |
images.append((seed, img)) | |
return images | |
# Function to modify image using ControlNet pipeline | |
def modify_image_with_controlnet(selected_image, additional_prompt, option='living_room', adj_1="spacious", architecture_style="modern", aesthetic="minimalist", | |
primary_color="brown", wood_finish="walnut", wall_color="beige", tiles="ceramic"): | |
gen_prompt = generate_prompt(option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles) | |
gen_prompt += '(('+additional_prompt+'))' | |
seed = selected_image[0] | |
generator = torch.manual_seed(seed) | |
cfg = 8 # Configuring guidance scale | |
steps = 30 # Number of inference steps | |
width, height = 640, 512 | |
# preprocessig for control net | |
img = np.array(selected_image[1]) | |
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) | |
edges = cv2.Canny(gray, 100, 200) | |
edges_image = Image.fromarray(edges) | |
controlnet_strength = 0.85 | |
image = controlnet_pipeline(prompt=gen_prompt, image=edges_image, guidance_scale=cfg, width=width, height=height, | |
controlnet_conditioning_scale=controlnet_strength, num_inference_steps=steps, generator=generator).images[0] | |
return image | |
# Gradio app logic | |
def generate_image(option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles, n_images, progress=gr.Progress()): | |
generated_images = generate_initial_image( | |
n_images=n_images, | |
option=option, | |
adj_1=adj_1, | |
architecture_style=architecture_style, | |
aesthetic=aesthetic, | |
primary_color=primary_color, | |
wood_finish=wood_finish, | |
wall_color=wall_color, | |
tiles=tiles | |
) | |
display_images = [img.resize((640, 512)) for _, img in generated_images] # Resize only for display | |
image_identifiers = [f"Image-{i+1}" for i in range(n_images)] | |
return display_images, generated_images, image_identifiers | |
def modify_image(selected_image_id, additional_prompt, option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles, images): | |
# Parse selected image identifier to get the index | |
image_index = int(selected_image_id.split("-")[1]) - 1 | |
selected_image = images[image_index] # Retrieve the (seed, original_img) tuple without resizing | |
# Modify image using ControlNet | |
modified_image = modify_image_with_controlnet( | |
selected_image=selected_image, | |
additional_prompt=additional_prompt, | |
option=option, | |
adj_1=adj_1, | |
architecture_style=architecture_style, | |
aesthetic=aesthetic, | |
primary_color=primary_color, | |
wood_finish=wood_finish, | |
wall_color=wall_color, | |
tiles=tiles | |
) | |
return modified_image | |
# Interface | |
with gr.Blocks() as app: | |
gr.Markdown("# Interior Design Image Generation and Modification") | |
# Image generation options | |
with gr.Row(): | |
option = gr.Radio(label="Room Type", choices=option_choices, value="living_room") | |
adj_1 = gr.Radio(label="Primary Adjective", choices=adj_1_choices, value="spacious") | |
architecture_style = gr.Radio(label="Architecture Style", choices=architecture_style_choices, value="modern") | |
with gr.Row(): | |
aesthetic = gr.Radio(label="Aesthetic", choices=aesthetic_choices, value="minimalist") | |
primary_color = gr.Radio(label="Primary Color", choices=primary_color_choices, value="neutral") | |
wood_finish = gr.Radio(label="Wood Finish", choices=wood_finish_choices, value="oak") | |
with gr.Row(): | |
wall_color = gr.Radio(label="Wall Color", choices=wall_color_choices, value="off-white") | |
tiles = gr.Radio(label="Tile Type", choices=tiles_choices, value="marble") | |
n_images = gr.Slider(label="Number of Images to Generate", minimum=1, maximum=5, step=1, value=1) | |
# Button to generate images | |
generate_btn = gr.Button("Generate Images") | |
output_images = gr.Gallery(label="Generated Images", columns=5) # Display resized images in gallery layout | |
output_generated_images = gr.State() # Hidden state to store original (unmodified) images data | |
output_identifiers = gr.State() # Hidden state for image identifiers | |
# Image selection and modification | |
selected_image_id = gr.Radio(label="Select an Image for Modification", choices=[]) | |
additional_prompt = gr.Textbox(label="Additional Prompt for Modification", placeholder="OPTIONAL") | |
modify_btn = gr.Button("Modify Image") | |
modified_image = gr.Image(label="Modified Image") | |
# Button actions | |
generate_btn.click( | |
fn=generate_image, | |
inputs=[option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles, n_images], | |
outputs=[output_images, output_generated_images, output_identifiers] | |
) | |
# Update radio button options after images are generated | |
output_images.change( | |
fn=lambda identifiers: gr.update(choices=identifiers, visible=True), # Populate with image identifiers | |
inputs=[output_identifiers], | |
outputs=[selected_image_id] | |
) | |
# Trigger modification on button click | |
modify_btn.click( | |
fn=modify_image, | |
inputs=[selected_image_id, additional_prompt, option, adj_1, architecture_style, aesthetic, primary_color, wood_finish, wall_color, tiles, output_generated_images], | |
outputs=modified_image | |
) | |
app.launch(debug=True, share=True) |