import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
from huggingface_hub import login
import os

# Authenticate with Hugging Face
login(token=os.environ.get('your_huggingface_token_here'))

# Load the Stable Fast 3D model
model_id = "stabilityai/stable-fast-3d"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

def convert_2d_to_3d(input_image, prompt):
    # Prepare the input image
    if input_image is not None:
        input_image = Image.open(io.BytesIO(input_image))
        input_image = input_image.resize((512, 512))
    
    # Generate the 3D preview
    output_image = pipe(
        prompt=prompt,
        image=input_image,
        num_inference_steps=50,
        guidance_scale=7.5
    ).images[0]
    
    return output_image

# Create the Gradio interface
iface = gr.Interface(
    fn=convert_2d_to_3d,
    inputs=[
        gr.Image(type="filepath", label="Upload 2D Floor Layout"),
        gr.Textbox(label="Prompt (e.g., '3D render of a modern apartment floor plan')")
    ],
    outputs=gr.Image(type="pil", label="3D Preview"),
    title="2D to 3D Floor Layout Converter",
    description="Upload a 2D floor layout image and get a 3D preview using Stable Fast 3D model."
)

# Launch the app
iface.launch()