feat: allow reference images for AI layer edits and relax stylistic constraints to allow for realism

backend/server.js

@@ -649,7 +649,7 @@ CRITICAL INSTRUCTIONS:
 
 // Edit Layer (Gemini Flash Image)
 app.post('/api/sparkplug/edit-layer', authenticate, async (req, res) => {
-  const { prompt, imageBase64, mimeType } = req.body;
+  const { prompt, imageBase64, mimeType, referenceImageBase64, referenceMimeType } = req.body;
   if (!prompt || !imageBase64) return res.status(400).json({ error: 'Missing prompt or image' });
 
   try {
@@ -660,18 +660,35 @@ app.post('/api/sparkplug/edit-layer', authenticate, async (req, res) => {
     const baseUrl = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-3.1-flash-image:generateContent';
     const base64Data = imageBase64.replace(/^data:image\/\w+;base64,/, "");
     
-    const payload = {
+    const parts = [
-      contents: [
+      { text: `You are an expert texture editor. Edit the following image based on this prompt: "${prompt}". 
-        {
+CRITICAL INSTRUCTIONS:
-          parts: [
+1. Output MUST be an image with the exact same aspect ratio as the input layer image.
-            { text: `You are an expert texture editor. Edit the following image based on this prompt: "${prompt}". \nCRITICAL INSTRUCTIONS:\n1. Output MUST be an image with the exact same aspect ratio as the input.\n2. Do NOT change the layout or mapping of the texture.\n3. Increase the resolution and quality of the visual elements as requested.` },
+2. Do NOT change the layout, mapping, or boundaries of the texture.
+3. You ARE encouraged to completely change the aesthetics, style, and realism of the texture as requested by the user, as long as the physical structure mapping remains intact.` },
       {
         inlineData: {
           mimeType: mimeType || "image/png",
           data: base64Data
         }
       }
-          ]
+    ];
+
+    if (referenceImageBase64) {
+      const refBase64Data = referenceImageBase64.replace(/^data:image\/\w+;base64,/, "");
+      parts.push({ text: "Use the following image purely as a stylistic reference for the edit:" });
+      parts.push({
+        inlineData: {
+          mimeType: referenceMimeType || "image/png",
+          data: refBase64Data
+        }
+      });
+    }
+    
+    const payload = {
+      contents: [
+        {
+          parts: parts
         }
       ],
       generationConfig: {

src/components/UVEditor.tsx

@@ -46,6 +46,7 @@ export function UVEditor({ onClose, onSave, targetObject, loadedModel, sceneObje
 
   const containerRef = useRef<HTMLDivElement>(null);
   const fileInputRef = useRef<HTMLInputElement>(null);
+  const refImageInputRef = useRef<HTMLInputElement>(null);
   const dragRef = useRef<{ handle: Handle; startX: number; startY: number; layer: Layer } | null>(null);
 
   const token = localStorage.getItem('auth_token') || '';
@@ -53,6 +54,8 @@ export function UVEditor({ onClose, onSave, targetObject, loadedModel, sceneObje
   const [isEditingLayer, setIsEditingLayer] = useState(false);
   const [aiPreviewUrl, setAiPreviewUrl] = useState<string | null>(null);
   const [originalLayerUrl, setOriginalLayerUrl] = useState<string | null>(null);
+  const [aiReferenceImage, setAiReferenceImage] = useState<File | null>(null);
+  const [aiReferencePreview, setAiReferencePreview] = useState<string | null>(null);
 
   const handleAiEdit = async (regenerate = false) => {
     if (!selectedId || !aiPrompt.trim()) return;
@@ -77,11 +80,40 @@ export function UVEditor({ onClose, onSave, targetObject, loadedModel, sceneObje
         // Use webp to preserve transparency while significantly reducing payload size
         const base64 = tempCanvas.toDataURL('image/webp', 0.8);
         
-        const fullPrompt = `The user has provided an image layer. Please apply the following edit exactly as requested, while strictly maintaining the EXACT same aspect ratio, resolution, and general structure of the original image. 
+        let referenceBase64 = undefined;
+        if (aiReferenceImage) {
+           const reader = new FileReader();
+           referenceBase64 = await new Promise<string>((resolve, reject) => {
+             reader.onload = (e) => {
+               const img = new Image();
+               img.onload = () => {
+                 const refCanvas = document.createElement('canvas');
+                 let rw = img.width;
+                 let rh = img.height;
+                 if (rw > 1024 || rh > 1024) {
+                   const rRatio = Math.min(1024 / rw, 1024 / rh);
+                   rw = Math.round(rw * rRatio);
+                   rh = Math.round(rh * rRatio);
+                 }
+                 refCanvas.width = rw;
+                 refCanvas.height = rh;
+                 const rctx = refCanvas.getContext('2d');
+                 rctx?.drawImage(img, 0, 0, rw, rh);
+                 resolve(refCanvas.toDataURL('image/webp', 0.8));
+               };
+               img.onerror = reject;
+               img.src = e.target?.result as string;
+             };
+             reader.onerror = reject;
+             reader.readAsDataURL(aiReferenceImage);
+           });
+        }
+        
+        const fullPrompt = `The user has provided an image layer. Please apply the following edit exactly as requested, while strictly maintaining the EXACT same aspect ratio and resolution of the original image. You may change the style, aesthetic, or realism. 
 User Prompt: "${aiPrompt}"
 Negative constraints: Do NOT change the aspect ratio. Do NOT add completely unrelated elements. Do NOT crop the image boundaries.`;
 
-        const newUrl = await editUVLayer(fullPrompt, base64, token || '');
+        const newUrl = await editUVLayer(fullPrompt, base64, token || '', referenceBase64);
         
         if (!originalLayerUrl && !regenerate) {
           setOriginalLayerUrl(layer.url);
@@ -108,6 +140,8 @@ Negative constraints: Do NOT change the aspect ratio. Do NOT add completely unre
     setAiPreviewUrl(null);
     setOriginalLayerUrl(null);
     setAiPrompt("");
+    setAiReferenceImage(null);
+    setAiReferencePreview(null);
   };
 
   const handleRejectEdit = () => {
@@ -118,11 +152,23 @@ Negative constraints: Do NOT change the aspect ratio. Do NOT add completely unre
         setLayers(ls => ls.map(l => l.id === selectedId ? { ...l, url: originalLayerUrl, img } : l));
         setAiPreviewUrl(null);
         setOriginalLayerUrl(null);
+        setAiReferenceImage(null);
+        setAiReferencePreview(null);
       };
       img.src = originalLayerUrl;
     } else {
       setAiPreviewUrl(null);
       setOriginalLayerUrl(null);
+      setAiReferenceImage(null);
+      setAiReferencePreview(null);
+    }
+  };
+
+  const handleRefImageUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const file = e.target.files?.[0];
+    if (file) {
+      setAiReferenceImage(file);
+      setAiReferencePreview(URL.createObjectURL(file));
     }
   };
 
@@ -526,6 +572,25 @@ Negative constraints: Do NOT change the aspect ratio. Do NOT add completely unre
                       </button>
                     </div>
                   ) : (
+                    <>
+                      <div className="flex items-center gap-2">
+                        <input type="file" ref={refImageInputRef} className="hidden" accept=".jpg,.jpeg,.png,.webp" onChange={handleRefImageUpload} />
+                        <button 
+                          onClick={() => refImageInputRef.current?.click()}
+                          className="flex-1 py-1.5 bg-zinc-800 hover:bg-zinc-700 rounded text-[10px] font-semibold text-zinc-300 border border-zinc-700 transition-colors flex items-center justify-center gap-1.5"
+                        >
+                          <Upload size={12} /> {aiReferencePreview ? 'Change Ref' : 'Optional Reference Image'}
+                        </button>
+                        {aiReferencePreview && (
+                          <div className="relative w-7 h-7 rounded overflow-hidden border border-zinc-600 flex-shrink-0 group">
+                            <img src={aiReferencePreview} alt="Ref" className="w-full h-full object-cover" />
+                            <div className="absolute inset-0 bg-black/60 flex items-center justify-center opacity-0 group-hover:opacity-100 cursor-pointer" 
+                                 onClick={() => { setAiReferenceImage(null); setAiReferencePreview(null); }}>
+                              <Trash2 size={10} className="text-red-400" />
+                            </div>
+                          </div>
+                        )}
+                      </div>
                       <button 
                         onClick={() => handleAiEdit(false)}
                         disabled={!aiPrompt.trim()} 
@@ -533,6 +598,7 @@ Negative constraints: Do NOT change the aspect ratio. Do NOT add completely unre
                       >
                         <span className="text-[14px] leading-none">✨</span> Generate Edit
                       </button>
+                    </>
                   )}
                 </div>
               </div>

src/lib/api/apiUtils.ts

@@ -80,7 +80,7 @@ export async function pollStatus<T>(
   throw new Error("Polling timeout exceeded.");
 }
 
-export async function editUVLayer(prompt: string, imageBase64: string, token: string): Promise<string> {
+export async function editUVLayer(prompt: string, imageBase64: string, token: string, referenceImageBase64?: string): Promise<string> {
   // Extract mime type if present in the data URI
   let mimeType = 'image/png';
   if (imageBase64.startsWith('data:')) {
@@ -90,13 +90,21 @@ export async function editUVLayer(prompt: string, imageBase64: string, token: st
     }
   }
 
+  let referenceMimeType;
+  if (referenceImageBase64 && referenceImageBase64.startsWith('data:')) {
+    const match = referenceImageBase64.match(/^data:([^;]+);/);
+    if (match) {
+      referenceMimeType = match[1];
+    }
+  }
+
   const response = await fetch('/api/sparkplug/edit-layer', {
     method: 'POST',
     headers: {
       'Content-Type': 'application/json',
       'Authorization': `Bearer ${token}`
     },
-    body: JSON.stringify({ prompt, imageBase64, mimeType })
+    body: JSON.stringify({ prompt, imageBase64, mimeType, referenceImageBase64, referenceMimeType })
   });
 
   if (!response.ok) {