Spaces:
Sleeping
Sleeping
File size: 3,737 Bytes
3d9c674 | 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 | import gradio as gr
import open3d as o3d
import matplotlib.pyplot as plt
import numpy as np
import os
from diffusers import DiffusionPipeline
import torch
from datasets import load_dataset
import transformers
import accelerate
# Load Stable Diffusion XL model
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipe.to("cuda" if torch.cuda.is_available() else "cpu")
# Ensure directories exist
os.makedirs("outputs/2d_drawings", exist_ok=True)
os.makedirs("outputs/3d_models", exist_ok=True)
# Parse plot size
def parse_plot_size(plot_size):
try:
dimensions = plot_size.lower().replace("feet", "").strip().split("x")
if len(dimensions) != 2:
raise ValueError("Invalid plot size format. Use '30x40 feet'.")
width, length = map(float, dimensions)
return width, length
except Exception as e:
raise ValueError(f"Error parsing plot size: {e}")
# Generate 2D drawing with Stable Diffusion
def generate_2d_drawing_stable_diffusion(plot_size, num_floors, storey_height, additional_reqs):
try:
# Create a descriptive prompt
width, length = parse_plot_size(plot_size)
height = num_floors * storey_height
prompt = (
f"Architectural floor plan of a {num_floors}-storey building with a plot size of "
f"{width}x{length} feet, each storey height {storey_height} feet. "
f"Additional details: {additional_reqs}"
)
# Generate image with Stable Diffusion
result = pipe(prompt)
generated_image = result.images[0]
# Save the generated image
output_path = os.path.join("outputs/2d_drawings", "2d_drawing_stable_diffusion.png")
generated_image.save(output_path)
return output_path
except Exception as e:
print(f"Error in generate_2d_drawing_stable_diffusion: {e}")
raise e
# Generate 3D model
def generate_3d_model(plot_size, num_floors, storey_height):
try:
width, length = parse_plot_size(plot_size)
height = num_floors * storey_height
mesh = o3d.geometry.TriangleMesh.create_box(width, length, height)
mesh.compute_vertex_normals()
output_path = os.path.join("outputs/3d_models", "3d_model.obj")
o3d.io.write_triangle_mesh(output_path, mesh)
return output_path
except Exception as e:
print(f"Error in generate_3d_model: {e}")
raise e
# Main function
def generate_building(plot_size, storey_height, num_floors, additional_reqs):
try:
# Generate 2D drawing using Stable Diffusion
drawing_path = generate_2d_drawing_stable_diffusion(plot_size, num_floors, storey_height, additional_reqs)
except Exception as e:
print(f"Stable Diffusion failed: {e}. Falling back to Matplotlib.")
drawing_path = generate_2d_drawing_matplotlib(plot_size, num_floors, storey_height)
# Generate 3D model
model_path = generate_3d_model(plot_size, num_floors, storey_height)
return drawing_path, model_path
# Gradio interface
iface = gr.Interface(
fn=generate_building,
inputs=[
gr.Textbox(label="Plot Size (e.g., 30x40 feet)", placeholder="Enter plot dimensions"),
gr.Number(label="Storey Height (ft)", value=10),
gr.Number(label="Number of Floors", value=3),
gr.Textbox(label="Additional Requirements", placeholder="Optional")
],
outputs=[
gr.Image(label="2D Drawing"),
gr.File(label="3D Model (Download)")
],
title="Building Design Generator with AI",
description="Generates realistic 2D drawings using Stable Diffusion XL and 3D models based on user inputs."
)
if __name__ == "__main__":
iface.launch()
|