Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 import torch
 import gc
 import spaces
@@ -11,173 +12,260 @@ import os
 torch.cuda.empty_cache()
 gc.collect()
 
-# Alpaca prompt template
-alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+# Model configuration
+BASE_MODEL = "Qwen/Qwen3-1.7B"
+ADAPTER_REPO = "vinoku89/svg-lora-sft-1.7b"
+ADAPTER_REVISION = "best"  # or specific tag like "qwen1.7b-sft-e2-14.2k-loss0.34-20260124-001"
 
-### Instruction:
-{}
+# Prompt configuration (from svg-sft-1.7b.yaml)
+SYSTEM_PROMPT = "You are an expert SVG code generator. Generate precise, well-formed SVG code that accurately matches the given description."
+USER_PREFIX = "Generate SVG code for the following description:\n\n"
+USER_SUFFIX = "\n\nProvide only the SVG code, starting with <svg and ending with </svg>."
 
-### Input:
-{}
+# Load model with LoRA adapter
+print(f"Loading base model: {BASE_MODEL}")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
 
-### Response:
-{}"""
-
-# Load model with memory optimizations
-model_path = "vinoku89/qwen3-4B-svg-code-gen"
-
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_path,
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
     torch_dtype=torch.float16,
     device_map="auto",
     low_cpu_mem_usage=True,
-    trust_remote_code=True  # Add this if needed for custom models
+    trust_remote_code=True
+)
+
+print(f"Loading LoRA adapter: {ADAPTER_REPO} (revision: {ADAPTER_REVISION})")
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER_REPO,
+    revision=ADAPTER_REVISION,
 )
+model.eval()
+print("Model loaded successfully!")
+
+
+def format_prompt(description: str) -> str:
+    """
+    Format the prompt using Qwen3 chat template.
+    """
+    user_content = f"{USER_PREFIX}{description}{USER_SUFFIX}"
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": user_content}
+    ]
+
+    # Use tokenizer's chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    return prompt
+
+
+def extract_svg(text: str) -> str:
+    """
+    Extract SVG code from generated text.
+    """
+    # Try to find complete SVG tags
+    svg_match = re.search(r'<svg[^>]*>.*?</svg>', text, re.DOTALL | re.IGNORECASE)
+    if svg_match:
+        return svg_match.group(0)
+
+    # If no complete SVG, try to find partial SVG and clean up
+    if '<svg' in text.lower():
+        # Find the start of SVG
+        start_idx = text.lower().find('<svg')
+        svg_content = text[start_idx:]
+
+        # If it doesn't end with </svg>, try to add it
+        if '</svg>' not in svg_content.lower():
+            # Find a good stopping point
+            svg_content = svg_content.split('<|')[0]  # Stop at chat tokens
+            svg_content = svg_content.split('\n\n')[0]  # Stop at double newline
+            if not svg_content.strip().endswith('</svg>'):
+                svg_content += '</svg>'
 
-def validate_svg(svg_content):
+        return svg_content
+
+    return text
+
+
+def validate_svg(svg_content: str):
     """
-    Validate if SVG content is properly formatted and renderable
+    Validate if SVG content is properly formatted and renderable.
     """
     try:
-        # Clean up the SVG content
         svg_content = svg_content.strip()
-        
+
         # If it doesn't start with <svg, try to extract SVG content
         if not svg_content.startswith('<svg'):
-            # Look for SVG tags in the content
-            svg_match = re.search(r'<svg[^>]*>.*?</svg>', svg_content, re.DOTALL | re.IGNORECASE)
-            if svg_match:
-                svg_content = svg_match.group(0)
-            else:
-                # If no complete SVG found, wrap content in SVG tags
-                if any(tag in svg_content.lower() for tag in ['<circle', '<rect', '<path', '<line', '<polygon', '<ellipse', '<text']):
-                    svg_content = f'<svg xmlns="http://www.w3.org/2000/svg" width="250" height="250">{svg_content}</svg>'
-                else:
-                    raise ValueError("No valid SVG elements found")
-        
+            svg_content = extract_svg(svg_content)
+
+        # Ensure xmlns is present
+        if 'xmlns=' not in svg_content:
+            svg_content = svg_content.replace('<svg', '<svg xmlns="http://www.w3.org/2000/svg"', 1)
+
         # Parse XML to validate structure
         ET.fromstring(svg_content)
-        
+
         return True, svg_content
-        
+
     except ET.ParseError as e:
         return False, f"XML Parse Error: {str(e)}"
     except Exception as e:
         return False, f"Validation Error: {str(e)}"
 
-@spaces.GPU(duration=60)  # Add duration limit
-def generate_svg(prompt):
+
+@spaces.GPU(duration=120)
+def generate_svg(description: str, temperature: float = 0.7, max_tokens: int = 2048):
+    """
+    Generate SVG code from a text description.
+    """
     # Clear cache before generation
     torch.cuda.empty_cache()
-    
-    # Format the prompt using Alpaca template
-    instruction = "Generate SVG code based on the given description."
-    formatted_prompt = alpaca_prompt.format(
-        instruction,
-        prompt,
-        ""  # Empty response - model will fill this
-    )
-    
-    inputs = tokenizer(formatted_prompt, return_tensors="pt")
-    
-    # Move inputs to the same device as model
+
+    # Format the prompt
+    prompt = format_prompt(description)
+
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    # Move inputs to model device
     if hasattr(model, 'device'):
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    
-    with torch.no_grad():  # Disable gradient computation to save memory
+    elif torch.cuda.is_available():
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+
+    # Generate
+    with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_length=1024,
-            temperature=0.7,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
             do_sample=True,
+            top_p=0.95,
             pad_token_id=tokenizer.eos_token_id,
-            max_new_tokens=512  # Limit new tokens instead of total length
+            eos_token_id=tokenizer.eos_token_id,
         )
-    
+
+    # Decode
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Extract only the response part (after "### Response:")
-    response_start = generated_text.find("### Response:")
-    if response_start != -1:
-        svg_code = generated_text[response_start + len("### Response:"):].strip()
+
+    # Extract SVG from the response
+    # The response should be after the assistant marker
+    if "<|im_start|>assistant" in generated_text:
+        svg_code = generated_text.split("<|im_start|>assistant")[-1]
+        svg_code = svg_code.replace("<|im_end|>", "").strip()
     else:
-        # Fallback: remove the original formatted prompt
-        svg_code = generated_text[len(formatted_prompt):].strip()
-    
+        # Fallback: extract everything after the prompt
+        svg_code = generated_text[len(prompt):].strip()
+
+    # Extract and clean SVG
+    svg_code = extract_svg(svg_code)
+
     # Validate SVG
     is_valid, result = validate_svg(svg_code)
-    
+
     if is_valid:
-        # SVG is valid
         validated_svg = result
-        # Ensure the SVG has proper dimensions for display (keep moderate size)
-        if 'width=' not in validated_svg or 'height=' not in validated_svg:
-            validated_svg = validated_svg.replace('<svg', '<svg width="250" height="250"', 1)
+        # Ensure reasonable display size
+        if 'width=' not in validated_svg.lower():
+            validated_svg = validated_svg.replace('<svg', '<svg width="400" height="400"', 1)
         svg_display = validated_svg
     else:
-        # SVG is invalid, show error message
         svg_display = f"""
-        <div style="width: 250px; height: 200px; border: 2px dashed #ff6b6b; 
-                    display: flex; align-items: center; justify-content: center; 
-                    background-color: #fff5f5; border-radius: 8px; padding: 15px; 
+        <div style="width: 400px; height: 300px; border: 2px dashed #ff6b6b;
+                    display: flex; align-items: center; justify-content: center;
+                    background-color: #fff5f5; border-radius: 8px; padding: 15px;
                     text-align: center; color: #e03131; font-family: Arial, sans-serif;">
             <div>
-                <h4 style="margin: 0 0 8px 0; color: #e03131;">🚫 Preview Not Available</h4>
+                <h4 style="margin: 0 0 8px 0; color: #e03131;">Preview Not Available</h4>
                 <p style="margin: 0; font-size: 12px;">Generated SVG contains errors:<br>
                 <em style="font-size: 11px;">{result}</em></p>
             </div>
         </div>
         """
-    
+
     # Clear cache after generation
     torch.cuda.empty_cache()
-    
+
     return svg_code, svg_display
 
-# Authentication function using HF Space secrets
-def authenticate(username, password):
-    """
-    Authentication function for Gradio using HF Space secrets
-    Returns True if credentials are valid, False otherwise
-    """
-    # Get credentials from HF Space secrets
-    valid_username = os.getenv("user")  # This matches your secret name "user"
-    valid_password = os.getenv("password")  # This matches your secret name "password"
-    
-    # Fallback credentials if secrets are not available (for local testing)
-    if valid_username is None:
-        valid_username = "user"
-        print("Warning: 'user' secret not found, using fallback")
-    
-    if valid_password is None:
-        valid_password = "password"
-        print("Warning: 'password' secret not found, using fallback")
-    
-    return username == valid_username and password == valid_password
-
-# Minimal CSS for slightly larger HTML preview only
+
+# Custom CSS
 custom_css = """
 div[data-testid="HTML"] {
-    min-height: 320px !important;
+    min-height: 420px !important;
+}
+.gradio-container {
+    max-width: 900px !important;
 }
 """
 
-gradio_app = gr.Interface(
-    fn=generate_svg,
-    inputs=gr.Textbox(
-        lines=2, 
-        placeholder="Describe the SVG you want (e.g., 'a red circle with blue border')..."
-    ),
-    outputs=[
-        gr.Code(label="Generated SVG Code", language="html"),
-        gr.HTML(label="SVG Preview")
-    ],
-    title="SVG Code Generator",
-    description="Generate SVG code from natural language using a fine-tuned LLM.",
-    css=custom_css
-)
+# Create Gradio interface
+with gr.Blocks(css=custom_css, title="SVG Code Generator") as gradio_app:
+    gr.Markdown("""
+    # SVG Code Generator
+
+    Generate SVG code from natural language descriptions using a fine-tuned Qwen3-1.7B model.
+
+    **Model:** `Qwen/Qwen3-1.7B` + LoRA adapter (`vinoku89/svg-lora-sft-1.7b`)
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            description_input = gr.Textbox(
+                label="Description",
+                placeholder="Describe the SVG you want (e.g., 'a red circle with blue border on white background')...",
+                lines=3
+            )
+
+            with gr.Row():
+                temperature = gr.Slider(
+                    minimum=0.1, maximum=1.5, value=0.7, step=0.1,
+                    label="Temperature"
+                )
+                max_tokens = gr.Slider(
+                    minimum=256, maximum=4096, value=2048, step=256,
+                    label="Max Tokens"
+                )
+
+            generate_btn = gr.Button("Generate SVG", variant="primary")
+
+        with gr.Column(scale=1):
+            svg_preview = gr.HTML(label="SVG Preview")
+
+    svg_code_output = gr.Code(label="Generated SVG Code", language="html", lines=15)
+
+    # Examples
+    gr.Examples(
+        examples=[
+            ["a simple red circle centered on a white background"],
+            ["a blue rectangle with rounded corners and a green border"],
+            ["a yellow star with 5 points"],
+            ["a gradient sunset with orange and purple colors"],
+            ["a simple house with a triangular roof and square windows"],
+        ],
+        inputs=description_input
+    )
+
+    # Connect the generate button
+    generate_btn.click(
+        fn=generate_svg,
+        inputs=[description_input, temperature, max_tokens],
+        outputs=[svg_code_output, svg_preview]
+    )
+
 
 if __name__ == "__main__":
-    gradio_app.launch(auth=(os.getenv("user"), os.getenv("password")), share=True, ssr_mode=False)
+    # Authentication using HF Space secrets
+    auth = None
+    user = os.getenv("user")
+    password = os.getenv("password")
+    if user and password:
+        auth = (user, password)
+
+    gradio_app.launch(auth=auth, share=True, ssr_mode=False)