Update app.py

app.py

@@ -6,17 +6,27 @@ import cadquery as cq
 from cadquery import exporters
 import tempfile
 import os
-from huggingface_hub import hf_hub_download
 import numpy as np
 
+# Set Hugging Face token (set in Space Secrets as HF_TOKEN)
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
 # Load 2D models
 model_id = "Osama03/Finetuned_diffusion_interiordesign"
-txt2img_pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-txt2img_pipe = txt2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-img2img_pipe = img2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+try:
+    txt2img_pipe = StableDiffusionPipeline.from_pretrained(
+        model_id, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
+    )
+    txt2img_pipe = txt2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+    img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        model_id, torch_dtype=torch.float16, use_auth_token=HF_TOKEN
+    )
+    img2img_pipe = img2img_pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+except Exception as e:
+    print(f"Failed to load {model_id}: {e}. Check token or model availability.")
+    raise
 
-# Try to load TripoSR (download weights directly)
+# Try to load TripoSR
 try:
     from tsr.system import TSR
     triposr_model = TSR.from_pretrained("stabilityai/TripoSR", config_name="config.yaml", weight_name="model.ckpt")
@@ -36,6 +46,74 @@ def generate_2d_image(prompt, style, features, input_image):
     
     return image
 
+def generate_3d_model(input_image, prompt):
+    if triposr_model is not None and input_image is not None:
+        try:
+            temp_image = Image.fromarray(input_image)
+            with tempfile.TemporaryDirectory() as tmpdir:
+                triposr_model(temp_image, device="cuda" if torch.cuda.is_available() else "cpu", export_format="glb", output_dir=tmpdir)
+                glb_path = os.path.join(tmpdir, "mesh.glb")
+                return glb_path
+        except Exception as e:
+            print(f"TripoSR failed: {e}. Using CadQuery fallback.")
+    
+    # Fallback: Simple CadQuery room
+    room = cq.Workplane("XY").box(5, 4, 3).faces(">Z").shell(-0.1)
+    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
+        exporters.export(room, tmpfile.name, exportType='GLB')
+        return tmpfile.name
+
+def generate_cad_model(room_length, room_width, room_height, features):
+    room = cq.Workplane("XY").box(room_length, room_width, room_height).faces(">Z").shell(-0.1)
+    feature_positions = {"sofa": (room_length/4, room_width/2, 0.5), "coffee table": (room_length/2, room_width/2, 0.3)}
+    feature_sizes = {"sofa": (2, 1, 0.8), "coffee table": (1, 1, 0.4)}
+    
+    for feat in features:
+        if feat in feature_positions:
+            pos = feature_positions[feat]
+            size = feature_sizes[feat]
+            room = room.union(cq.Workplane("XY").transformed(offset=(pos[0], pos[1], pos[2])).box(*size))
+    
+    with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile:
+        exporters.export(room, tmpfile.name, exportType='GLB')
+        return tmpfile.name
+
+# UI with Tabs
+with gr.Blocks(title="Interior Design Generator MVP") as demo:
+    gr.Markdown("# Interior Design Image Generator MVP with 3D & CAD")
+    
+    with gr.Tabs():
+        with gr.Tab("2D Generation"):
+            with gr.Row():
+                prompt = gr.Textbox(label="Text Prompt")
+                style = gr.Dropdown(choices=["modern", "vintage", "minimalist", "industrial", "bohemian", "scandinavian", "rustic"], value="modern")
+                features = gr.Checkboxgroup(choices=["sofa", "coffee table", "lamp", "bookshelf", "fireplace", "plants", "artwork", "rug"])
+            input_image = gr.Image(label="Upload Input Image (optional)", type="numpy")
+            generate_2d_btn = gr.Button("Generate 2D Image")
+            output_2d = gr.Image(label="Generated 2D Design")
+            generate_2d_btn.click(generate_2d_image, inputs=[prompt, style, features, input_image], outputs=output_2d)
+        
+        with gr.Tab("3D Visualization"):
+            gr.Markdown("Generate and view 3D from the 2D image (or prompt if no image).")
+            input_3d_image = gr.Image(label="Use Generated 2D Image (paste from above or upload)", type="numpy")
+            input_3d_prompt = gr.Textbox(label="Fallback Prompt (if no image)")
+            generate_3d_btn = gr.Button("Generate 3D Model")
+            output_3d = gr.Model3D(label="3D Model (rotate/zoom in browser)")
+            generate_3d_btn.click(generate_3d_model, inputs=[input_3d_image, input_3d_prompt], outputs=output_3d)
+        
+        with gr.Tab("CAD-like Editing"):
+            gr.Markdown("Parametric editing: Adjust room dimensions and features, regenerate CAD model.")
+            room_length = gr.Slider(3, 10, value=5, label="Room Length (m)")
+            room_width = gr.Slider(3, 10, value=4, label="Room Width (m)")
+            room_height = gr.Slider(2, 4, value=3, label="Room Height (m)")
+            cad_features = gr.Checkboxgroup(choices=["sofa", "coffee table"], label="Features (positioned simply)")
+            generate_cad_btn = gr.Button("Generate CAD Model")
+            output_cad = gr.Model3D(label="CAD 3D Model (edit params to update)")
+            generate_cad_btn.click(generate_cad_model, inputs=[room_length, room_width, room_height, cad_features], outputs=output_cad)
+
+if __name__ == "__main__":
+    demo.launch()    return image
+
 def generate_3d_model(input_image, prompt):
     if triposr_model is not None and input_image is not None:
         # Use TripoSR for 3D