app.py

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)