Spaces:
Runtime error
Runtime error
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) |