Spaces:
Sleeping
Sleeping
| 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() | |