Create app.py

app.py

@@ -0,0 +1,211 @@
+import gradio as gr
+import os
+import torch
+import cadquery as cq
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoConfig
+from PIL import Image
+import ast  # For safe evaluation of string-formatted lists
+from io import BytesIO
+
+# --- CONFIGURATION (Keep as constants) ---
+MODEL_PATH = "/raid/home/posahemanth/miniconda3/Sai/FinalYearProject/1000_gpusoutput"
+OUTPUT_DIRECTORY = "/raid/home/posahemanth/miniconda3/Sai/FinalYearProject/Gradio_Output"  # Separate output
+USE_FLASH_ATTENTION = True
+PRE_TRAINED_MODEL_NAME = "microsoft/Phi-4-multimodal-instruct"
+os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)  # Ensure the output directory exists
+
+
+# --- MODEL LOADING (Global Scope) ---
+# Load only once, outside the functions, to improve performance
+try:
+    config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True, local_files_only=True)
+    config.attn_implementation = "flash_attention_2" if USE_FLASH_ATTENTION else "sdpa"
+    config.num_logits_to_keep = 20
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        config=config,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16 if USE_FLASH_ATTENTION else torch.float32,
+        local_files_only=True
+    ).to("cuda").eval()  # .eval() is crucial for inference
+
+    processor = AutoProcessor.from_pretrained(
+        PRE_TRAINED_MODEL_NAME,
+        trust_remote_code=True,
+        local_files_only=False,
+        config=config,
+    )
+except Exception as e:
+    print(f"Error loading model/processor: {e}")
+    raise  # Re-raise to halt execution
+
+
+# --- CAPTION GENERATION ---
+def generate_caption(image):
+    """Generates a caption for the given image."""
+    if image is None:
+        return "Please upload an image."
+
+    try:
+        # Convert numpy array to PIL Image
+        image = Image.fromarray(image).convert("RGB")
+    except Exception as e:
+        print(f"Error converting image: {e}")
+        return "Error processing image."
+    
+    prompt = "Describe this image."
+    user_message = {'role': 'user', 'content': f'<|image_1|>{prompt}'}
+    prompt_tokenized = processor.tokenizer.apply_chat_template([user_message], tokenize=False, add_generation_prompt=True)
+    inputs = processor(prompt_tokenized, images=[image], return_tensors='pt').to("cuda")
+
+    try:
+        with torch.no_grad(): # Ensure no gradients are calculated
+            generated_ids = model.generate(
+                **inputs,
+                eos_token_id=processor.tokenizer.eos_token_id,
+                max_new_tokens=512,
+                num_logits_to_keep=20,
+            )
+
+        input_len = inputs.input_ids.size(1)
+        generated_text = processor.decode(
+            generated_ids[0, input_len:],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        ).strip()
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        return "Error during caption generation."
+
+    return generated_text
+
+# --- CAD MODEL BUILDING ---
+def build_model(sequence):
+    """Builds a CAD model from the sequence and returns the STEP file path."""
+    workplane = cq.Workplane("XY")
+    model = None
+    primitive = None
+
+    if isinstance(sequence, str):
+        try:
+            sequence = ast.literal_eval(sequence)
+        except (ValueError, SyntaxError):
+            return "Invalid sequence format.  Could not convert to list."
+        if not isinstance(sequence, list):
+            return "Invalid sequence format.  Expected a list."
+    elif not isinstance(sequence, list):
+        return "Invalid sequence format.  Expected a list."
+
+
+    for step in sequence:
+        index = step[0]
+        if index == 0:  # Cube
+            _, length, width, height, loc_x, loc_y, loc_z, axis = step
+            primitive = workplane.box(length, width, height).translate((loc_x, loc_y, loc_z))
+        elif index == 1:  # Cylinder
+            _, height, radius, loc_x, loc_y, loc_z, axis = step
+            primitive = workplane.cylinder(radius, height).translate((loc_x, loc_y, loc_z))
+        elif index == 2:  # Sphere
+            _, radius, loc_x, loc_y, loc_z, axis = step
+            primitive = workplane.sphere(radius).translate((loc_x, loc_y, loc_z))
+
+        if primitive is None:
+            print(f"Skipping step {step} because primitive was not initialized.")
+            continue
+
+        if index in [3, 4, 5]:  # Operations
+            if model is None:
+                model = primitive
+            _, loc_x, loc_y, loc_z = step
+            if index == 3:
+                model = model.union(primitive.translate((loc_x, loc_y, loc_z)))
+            elif index == 4:
+                model = model.cut(primitive.translate((loc_x, loc_y, loc_z)))
+            elif index == 5:
+                model = model.intersect(primitive.translate((loc_x, loc_y, loc_z)))
+
+        if model is None:
+            model = primitive
+
+    if model is None:
+        return "Error: No valid CAD model was created."
+
+    # Create a unique filename using a timestamp (more robust)
+    import datetime
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    model_name = f"generated_model_{timestamp}"
+    step_file_path = os.path.join(OUTPUT_DIRECTORY, f"{model_name}.step")
+    cq.exporters.export(model, step_file_path)
+    return step_file_path
+
+
+
+
+def process_image(image):
+    """Combines caption generation and model building."""
+    if image is None:
+        return "Please upload an image first.", None
+
+    caption = generate_caption(image)
+    if not caption or caption.startswith("Error"):
+        return caption, None
+
+    step_file_path = build_model(caption)
+    if step_file_path.startswith("Error"):
+        return step_file_path, None
+
+    return "CAD model generated successfully!", step_file_path
+
+
+
+# --- GRADIO INTERFACE ---
+
+css = """
+.container {
+    max-width: 800px;
+    margin: auto;
+    padding: 20px;
+    border: 2px solid #ddd;
+    border-radius: 10px;
+}
+h1 {
+    text-align: center;
+    color: #333;
+}
+.description {
+    text-align: center;
+    margin-bottom: 20px;
+}
+.input-section, .output-section {
+    margin-bottom: 20px;
+    padding: 10px;
+    border: 1px solid #ccc;
+    border-radius: 5px;
+}
+.input-section h2, .output-section h2 {
+    margin-top: 0;
+    color: #555;
+}
+.output-section p {
+    font-weight: bold;
+}
+
+"""
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(label="Upload Image", type="numpy"),
+    outputs=[
+        gr.Textbox(label="Status"),  # Show status messages
+        gr.File(label="Download STEP File")  # Download link for the file
+    ],
+    title="Image to CAD Converter",
+    description="Upload an image of a mechanical drawing, and this app will attempt to generate a corresponding STEP CAD file.",
+    css=css, # Apply the CSS
+    allow_flagging="never",  # Disable flagging
+    theme=gr.themes.Soft()
+
+)
+
+iface.launch(share=True)