Create app.py

app.py

@@ -0,0 +1,256 @@
+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")
+
+# 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)