Update app.py

app.py

@@ -1,19 +1,29 @@
 import gradio as gr
 from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
+from peft import PeftModel
 import torch
 from PIL import Image
 
-# YAHAN APNA MODEL ID LAZMI CHANGE KAREIN (Jo HF par upload kiya tha)
-model_id = "ocilab/flowdex-sketch-model"
+# 1. Aapka apna model (Jisme naya knowledge hai)
+adapter_id = "ocilab/flowdex-sketch-model"
+# 2. Original Base Model (Jiske upar knowledge attach hoga)
+base_model_id = "unsloth/Qwen2-VL-2B-Instruct"
 
-print("Model CPU memory mein load ho raha hai... (Thoda wait karein)")
-processor = AutoProcessor.from_pretrained(model_id)
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    model_id, 
+print("AI Pipeline Setup ho rahi hai... (Takes a moment on CPU)")
+
+# Processor aapke adapter se hi load hoga
+processor = AutoProcessor.from_pretrained(adapter_id)
+
+# Pehle Base Model load karein
+base_model = Qwen2VLForConditionalGeneration.from_pretrained(
+    base_model_id, 
     torch_dtype=torch.bfloat16, 
-    device_map="cpu" # Free space ke liye CPU lazmi hai
+    device_map="cpu"
 )
 
+# Ab base model ke upar aapki training attach karein
+model = PeftModel.from_pretrained(base_model, adapter_id)
+
 def generate_ui(image):
     if image is None:
         return "Please upload a sketch image first."
@@ -41,7 +51,6 @@ def generate_ui(image):
     
     return generated_code
 
-# Gradio ka Interface jo frontend testing UI aur API dono banayega
 iface = gr.Interface(
     fn=generate_ui,
     inputs=gr.Image(type="pil", label="Upload Sketch"),