ShihadShowkat/ui_to_code

1import io2import os3import re4import traceback5 6import torch7from fastapi import FastAPI, UploadFile, File8from fastapi.middleware.cors import CORSMiddleware9from fastapi.responses import JSONResponse10from PIL import Image11from transformers import AutoProcessor, AutoModelForImageTextToText12 13# ==========================================================14# FastAPI15# ==========================================================16 17app = FastAPI(18    title="UI to Code API",19    description="Generate HTML, CSS and React JSX from UI screenshots."20)21 22app.add_middleware(23    CORSMiddleware,24    allow_origins=["*"],25    allow_credentials=True,26    allow_methods=["*"],27    allow_headers=["*"],28)29 30# ==========================================================31# Model Configuration32# ==========================================================33 34MODEL_ID = "ShihadShowkat/unsloth_finetuned_V2"35HF_TOKEN = os.getenv("HF_TOKEN")36 37DEVICE = "cpu"38 39print("=" * 60)40print("Loading model...")41print(f"Device: {DEVICE}")42print("=" * 60)43 44processor = AutoProcessor.from_pretrained(45    MODEL_ID,46    token=HF_TOKEN47)48 49model = AutoModelForImageTextToText.from_pretrained(50    MODEL_ID,51    torch_dtype=torch.float32,52    device_map=DEVICE,53    token=HF_TOKEN54)55 56model.eval()57 58print("Model loaded successfully.")59print("=" * 60)60 61# ==========================================================62# Prompt63# ==========================================================64 65PROMPT = (66    """67        Generate the HTML, CSS and React JSX implementation of this UI.68        69        Return ONLY the following markdown code blocks:70        71        ```html72        ```73        74        ```css75        ```76        77        ```jsx78        ```79    """80)81 82# ==========================================================83# Helper Function84# ==========================================================85 86def extract_code_block(text: str, language: str) -> str:87    """88    Extract a markdown code block by language.89    """90 91    pattern = rf"```{language}\s*(.*?)```"92    match = re.search(pattern, text, flags=re.DOTALL | re.IGNORECASE)93 94    if match:95        return match.group(1).strip()96 97    return ""98 99# ==========================================================100# Health Check101# ==========================================================102 103@app.get("/")104def health_check():105    return {106        "status": "online",107        "model": MODEL_ID,108        "device": DEVICE109    }110 111# ==========================================================112# Generate Endpoint113# ==========================================================114 115@app.post("/generate-code")116async def generate_code(image: UploadFile = File(...)):117 118    try:119 120        # -----------------------------------------121        # Read Image122        # -----------------------------------------123 124        image_bytes = await image.read()125 126        pil_image = Image.open(127            io.BytesIO(image_bytes)128        ).convert("RGB")129 130        # Resize large images to improve inference speed131        pil_image.thumbnail((768, 768))132 133        # -----------------------------------------134        # Build Conversation135        # -----------------------------------------136 137        messages = [138            {139                "role": "user",140                "content": [141                    {142                        "type": "text",143                        "text": PROMPT144                    },145                    {146                        "type": "image",147                        "image": pil_image148                    }149                ]150            }151        ]152 153        # -----------------------------------------154        # Tokenization155        # -----------------------------------------156 157        prompt = processor.apply_chat_template(158            messages,159            add_generation_prompt=True160        )161 162        inputs = processor(163            text=prompt,164            images=pil_image,165            return_tensors="pt"166        )167 168        inputs = inputs.to(DEVICE)169 170        # -----------------------------------------171        # Inference172        # -----------------------------------------173 174        with torch.inference_mode():175 176            outputs = model.generate(177                **inputs,178                max_new_tokens=1200,179                do_sample=False,180                use_cache=True,181                eos_token_id=processor.tokenizer.eos_token_id,182                pad_token_id=processor.tokenizer.eos_token_id,183            )184 185        generated_tokens = outputs[0][len(inputs.input_ids[0]):]186 187        generated_text = processor.decode(188            generated_tokens,189            skip_special_tokens=True190        )191 192        # -----------------------------------------193        # Extract HTML / CSS / JSX194        # -----------------------------------------195 196        html = extract_code_block(generated_text, "html")197        css = extract_code_block(generated_text, "css")198        jsx = extract_code_block(generated_text, "jsx")199 200        # -----------------------------------------201        # Return JSON202        # -----------------------------------------203 204        return JSONResponse(205            content={206                "success": True,207                "html": html,208                "css": css,209                "jsx": jsx,210                "raw": generated_text211            }212        )213 214    except Exception as e:215 216        traceback.print_exc()217 218        return JSONResponse(219            status_code=500,220            content={221                "success": False,222                "error": str(e)223            }224        )