浏览代码

fix: resolve webui freezing issues with proper cleanup and error handling

- Add polling cleanup with isPolling flags to prevent infinite loops
- Implement error boundary for graceful error handling
- Add global error handlers for unhandled rejections
- Create safe localStorage hooks that exclude large image data
- Separate image storage from form data to prevent quota issues
- Fix TypeScript compilation errors in storage and image handling
- Add memory management utilities for better resource cleanup

Resolves webui freezing after extended use due to memory leaks and infinite polling.
Fszontagh 3 月之前
父节点
当前提交
05be52a612

+ 69 - 25
webui/app/img2img/page.tsx

@@ -1,19 +1,18 @@
 'use client';
 
 import { useState, useRef, useEffect } from 'react';
-import { Header } from '@/components/header';
-import { AppLayout } from '@/components/layout';
+import { Header, AppLayout } from '@/components/layout';
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Textarea } from '@/components/ui/textarea';
-import { PromptTextarea } from '@/components/prompt-textarea';
+import { PromptTextarea } from '@/components/forms';
 import { Label } from '@/components/ui/label';
 import { Card, CardContent } from '@/components/ui/card';
 import { ImageInput } from '@/components/ui/image-input';
 import { apiClient, type JobInfo } from '@/lib/api';
 import { Loader2, Download, X } from 'lucide-react';
 import { downloadImage, downloadAuthenticatedImage, fileToBase64 } from '@/lib/utils';
-import { useLocalStorage } from '@/lib/hooks';
+import { useLocalStorage, useMemoryStorage } from '@/lib/storage';
 
 type Img2ImgFormData = {
   prompt: string;
@@ -43,10 +42,18 @@ const defaultFormData: Img2ImgFormData = {
 
 function Img2ImgForm() {
 
-  const [formData, setFormData] = useLocalStorage<Img2ImgFormData>(
+  // Store form data without the image to avoid localStorage quota issues
+  const { image: _, ...formDataWithoutImage } = defaultFormData;
+  const [formData, setFormData] = useLocalStorage<Omit<Img2ImgFormData, 'image'>>(
     'img2img-form-data',
-    defaultFormData
+    formDataWithoutImage
   );
+  
+  // Store image separately in memory
+  const [imageData, setImageData] = useMemoryStorage<string>('');
+  
+  // Combined form data with image
+  const fullFormData = { ...formData, image: imageData };
 
   const [loading, setLoading] = useState(false);
   const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
@@ -59,6 +66,16 @@ function Img2ImgForm() {
   const [originalImage, setOriginalImage] = useState<string | null>(null);
   const [isResizing, setIsResizing] = useState(false);
   const [error, setError] = useState<string | null>(null);
+  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
+
+  // Cleanup polling on unmount
+  useEffect(() => {
+    return () => {
+      if (pollCleanup) {
+        pollCleanup();
+      }
+    };
+  }, [pollCleanup]);
 
   useEffect(() => {
     const loadModels = async () => {
@@ -94,6 +111,7 @@ function Img2ImgForm() {
     setError(null);
 
     if (!image) {
+      setImageData('');
       setFormData(prev => ({ ...prev, image: '' }));
       setPreviewImage(null);
       setImageValidation(null);
@@ -117,7 +135,8 @@ function Img2ImgForm() {
 
       // Store original image for resizing
       setOriginalImage(imageBase64);
-      setFormData(prev => ({ ...prev, image: imageBase64 }));
+      setImageData(imageBase64);
+      setFormData(prev => ({ ...prev })); // Don't store image in localStorage
       setPreviewImage(previewUrl);
     } catch (err) {
       setError('Failed to process image');
@@ -139,12 +158,20 @@ function Img2ImgForm() {
 
       try {
         setIsResizing(true);
+        
+        // Validate image data before sending
+        if (!originalImage.startsWith('data:image/')) {
+          console.warn('Invalid image data for resizing');
+          return;
+        }
+
         const result = await apiClient.resizeImage(originalImage, formData.width, formData.height);
-        setFormData(prev => ({ ...prev, image: result.image }));
+        setImageData(result.image);
+        setFormData(prev => ({ ...prev })); // Don't store image in localStorage
         setPreviewImage(result.image);
       } catch (err) {
         console.error('Failed to resize image:', err);
-        setError('Failed to resize image');
+        setError('Failed to resize image - you may need to resize manually');
       } finally {
         setIsResizing(false);
       }
@@ -159,14 +186,23 @@ function Img2ImgForm() {
       setError(result.error || 'Invalid image');
     } else {
       setError(null);
+      // If we have base64 data from URL download, use it for preview
+      if (result.base64Data && selectedImage && typeof selectedImage === 'string') {
+        setPreviewImage(result.base64Data);
+        setFormData(prev => ({ ...prev, image: result.base64Data }));
+        setOriginalImage(result.base64Data);
+      }
     }
   };
 
   const pollJobStatus = async (jobId: string) => {
     const maxAttempts = 300;
     let attempts = 0;
+    let isPolling = true;
 
     const poll = async () => {
+      if (!isPolling) return;
+      
       try {
         const status = await apiClient.getJobStatus(jobId);
         setJobInfo(status);
@@ -200,32 +236,44 @@ function Img2ImgForm() {
           // Create a new array to trigger React re-render
           setGeneratedImages([...imageUrls]);
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'failed') {
           setError(status.error || 'Generation failed');
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'cancelled') {
           setError('Generation was cancelled');
           setLoading(false);
+          isPolling = false;
         } else if (attempts < maxAttempts) {
           attempts++;
           setTimeout(poll, 2000);
         } else {
           setError('Job polling timeout');
           setLoading(false);
+          isPolling = false;
         }
       } catch (err) {
-        setError(err instanceof Error ? err.message : 'Failed to check job status');
-        setLoading(false);
+        if (isPolling) {
+          setError(err instanceof Error ? err.message : 'Failed to check job status');
+          setLoading(false);
+          isPolling = false;
+        }
       }
     };
 
     poll();
+
+    // Return cleanup function
+    return () => {
+      isPolling = false;
+    };
   };
 
   const handleGenerate = async (e: React.FormEvent) => {
     e.preventDefault();
 
-    if (!formData.image) {
+    if (!fullFormData.image) {
       setError('Please upload or select an image first');
       return;
     }
@@ -243,14 +291,15 @@ function Img2ImgForm() {
 
     try {
       const requestData = {
-        ...formData,
+        ...fullFormData,
       };
 
       const job = await apiClient.img2img(requestData);
       setJobInfo(job);
       const jobId = job.request_id || job.id;
       if (jobId) {
-        await pollJobStatus(jobId);
+        const cleanup = pollJobStatus(jobId);
+        setPollCleanup(() => cleanup);
       } else {
         setError('No job ID returned from server');
         setLoading(false);
@@ -268,6 +317,11 @@ function Img2ImgForm() {
         await apiClient.cancelJob(jobId);
         setLoading(false);
         setError('Generation cancelled');
+        // Cleanup polling
+        if (pollCleanup) {
+          pollCleanup();
+          setPollCleanup(null);
+        }
       } catch (err) {
         console.error('Failed to cancel job:', err);
       }
@@ -441,7 +495,7 @@ function Img2ImgForm() {
 
 
                 <div className="flex gap-2">
-                  <Button type="submit" disabled={loading || !formData.image || (imageValidation && !imageValidation.isValid)} className="flex-1">
+                  <Button type="submit" disabled={loading || !imageData || (imageValidation && !imageValidation.isValid)} className="flex-1">
                     {loading ? (
                       <>
                         <Loader2 className="h-4 w-4 animate-spin" />
@@ -464,16 +518,6 @@ function Img2ImgForm() {
                     {error}
                   </div>
                 )}
-
-                {loading && jobInfo && (
-                  <div className="rounded-md bg-muted p-3 text-sm">
-                    <p>Job ID: {jobInfo.id || jobInfo.request_id || 'N/A'}</p>
-                    <p>Status: {jobInfo.status}</p>
-                    {jobInfo.progress !== undefined && (
-                      <p>Progress: {Math.round(jobInfo.progress * 100)}%</p>
-                    )}
-                  </div>
-                )}
               </form>
             </CardContent>
           </Card>

+ 61 - 50
webui/app/inpainting/page.tsx

@@ -1,28 +1,26 @@
 'use client';
 
 import { useState, useEffect } from 'react';
-import { Header } from '@/components/header';
+import { Header } from '@/components/layout';
 import { AppLayout } from '@/components/layout';
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Textarea } from '@/components/ui/textarea';
-import { PromptTextarea } from '@/components/prompt-textarea';
+import { PromptTextarea } from '@/components/forms';
 import { Label } from '@/components/ui/label';
-import { Card, CardContent } from '@/components/ui/card';
-import { InpaintingCanvas } from '@/components/inpainting-canvas';
+import { Card, CardContent, CardHeader, CardTitle, CardDescription } from '@/components/ui/card';
+import { InpaintingCanvas } from '@/components/features/image-generation';
 import { apiClient, type JobInfo } from '@/lib/api';
 import { Loader2, X, Download } from 'lucide-react';
 import { downloadAuthenticatedImage } from '@/lib/utils';
-import { useLocalStorage } from '@/lib/hooks';
+import { useLocalStorage } from '@/lib/storage';
 import { ModelSelectionProvider, useModelSelection, useCheckpointSelection, useModelTypeSelection } from '@/contexts/model-selection-context';
-import { EnhancedModelSelect, EnhancedModelSelectGroup } from '@/components/enhanced-model-select';
-import { ModelSelectionWarning, AutoSelectionStatus } from '@/components/model-selection-indicator';
+import { EnhancedModelSelect } from '@/components/features/models';
+// import { AutoSelectionStatus } from '@/components/features/models';
 
 type InpaintingFormData = {
   prompt: string;
   negative_prompt: string;
-  source_image: string;
-  mask_image: string;
   steps: number;
   cfg_scale: number;
   seed: string;
@@ -35,8 +33,6 @@ type InpaintingFormData = {
 const defaultFormData: InpaintingFormData = {
   prompt: '',
   negative_prompt: '',
-  source_image: '',
-  mask_image: '',
   steps: 20,
   cfg_scale: 7.5,
   seed: '',
@@ -70,15 +66,30 @@ function InpaintingForm() {
 
   const [formData, setFormData] = useLocalStorage<InpaintingFormData>(
     'inpainting-form-data',
-    defaultFormData
+    defaultFormData,
+    { excludeLargeData: true, maxSize: 512 * 1024 }
   );
 
+  // Separate state for image data (not stored in localStorage)
+  const [sourceImage, setSourceImage] = useState<string>('');
+  const [maskImage, setMaskImage] = useState<string>('');
+
   const [loading, setLoading] = useState(false);
   const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
   const [generatedImages, setGeneratedImages] = useState<string[]>([]);
   const [loraModels, setLoraModels] = useState<string[]>([]);
   const [embeddings, setEmbeddings] = useState<string[]>([]);
   const [error, setError] = useState<string | null>(null);
+  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
+
+  // Cleanup polling on unmount
+  useEffect(() => {
+    return () => {
+      if (pollCleanup) {
+        pollCleanup();
+      }
+    };
+  }, [pollCleanup]);
 
   useEffect(() => {
     const loadModels = async () => {
@@ -121,20 +132,23 @@ function InpaintingForm() {
   };
 
   const handleSourceImageChange = (image: string) => {
-    setFormData((prev) => ({ ...prev, source_image: image }));
+    setSourceImage(image);
     setError(null);
   };
 
   const handleMaskImageChange = (image: string) => {
-    setFormData((prev) => ({ ...prev, mask_image: image }));
+    setMaskImage(image);
     setError(null);
   };
 
   const pollJobStatus = async (jobId: string) => {
-    const maxAttempts = 300;
+    const maxAttempts = 300; // 5 minutes with 2 second interval
     let attempts = 0;
+    let isPolling = true;
 
     const poll = async () => {
+      if (!isPolling) return;
+      
       try {
         const status = await apiClient.getJobStatus(jobId);
         setJobInfo(status);
@@ -168,37 +182,49 @@ function InpaintingForm() {
           // Create a new array to trigger React re-render
           setGeneratedImages([...imageUrls]);
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'failed') {
           setError(status.error || 'Generation failed');
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'cancelled') {
           setError('Generation was cancelled');
           setLoading(false);
+          isPolling = false;
         } else if (attempts < maxAttempts) {
           attempts++;
           setTimeout(poll, 2000);
         } else {
           setError('Job polling timeout');
           setLoading(false);
+          isPolling = false;
         }
       } catch (err) {
-        setError(err instanceof Error ? err.message : 'Failed to check job status');
-        setLoading(false);
+        if (isPolling) {
+          setError(err instanceof Error ? err.message : 'Failed to check job status');
+          setLoading(false);
+          isPolling = false;
+        }
       }
     };
 
     poll();
+
+    // Return cleanup function
+    return () => {
+      isPolling = false;
+    };
   };
 
   const handleGenerate = async (e: React.FormEvent) => {
     e.preventDefault();
 
-    if (!formData.source_image) {
+    if (!sourceImage) {
       setError('Please upload a source image first');
       return;
     }
 
-    if (!formData.mask_image) {
+    if (!maskImage) {
       setError('Please create a mask first');
       return;
     }
@@ -221,6 +247,8 @@ function InpaintingForm() {
 
       const requestData = {
         ...formData,
+        source_image: sourceImage,
+        mask_image: maskImage,
         model: selectedCheckpoint || undefined,
         vae: selectedVae || undefined,
       };
@@ -229,7 +257,8 @@ function InpaintingForm() {
       setJobInfo(job);
       const jobId = job.request_id || job.id;
       if (jobId) {
-        await pollJobStatus(jobId);
+        const cleanup = pollJobStatus(jobId);
+        setPollCleanup(() => cleanup);
       } else {
         setError('No job ID returned from server');
         setLoading(false);
@@ -247,6 +276,11 @@ function InpaintingForm() {
         await apiClient.cancelJob(jobId);
         setLoading(false);
         setError('Generation cancelled');
+        // Cleanup polling
+        if (pollCleanup) {
+          pollCleanup();
+          setPollCleanup(null);
+        }
       } catch (err) {
         console.error('Failed to cancel job:', err);
       }
@@ -405,10 +439,11 @@ function InpaintingForm() {
                   </div>
   
                   {/* Model Selection Section */}
-                  <EnhancedModelSelectGroup
-                    title="Model Selection"
-                    description="Select the checkpoint and additional models for generation"
-                  >
+                  <Card>
+                    <CardHeader>
+                      <CardTitle>Model Selection</CardTitle>
+                      <CardDescription>Select checkpoint and additional models for generation</CardDescription>
+                    </CardHeader>
                     {/* Checkpoint Selection */}
                     <div className="space-y-2">
                       <Label htmlFor="checkpoint">Checkpoint Model *</Label>
@@ -444,26 +479,12 @@ function InpaintingForm() {
                       placeholder="Use default VAE"
                     />
   
-                    {/* Auto-selection Status */}
-                    <div className="pt-2">
-                      <AutoSelectionStatus
-                        isAutoSelecting={isAutoSelecting}
-                        hasAutoSelection={Object.keys(state.autoSelectedModels).length > 0}
-                      />
-                    </div>
-  
-                    {/* Warnings and Errors */}
-                    <ModelSelectionWarning
-                      warnings={warnings}
-                      errors={error ? [error] : []}
-                      onClearWarnings={actions.clearWarnings}
-                    />
-                  </EnhancedModelSelectGroup>
+                  </Card>
   
                   <div className="flex gap-2">
                     <Button
                       type="submit"
-                      disabled={loading || !formData.source_image || !formData.mask_image}
+                      disabled={loading || !sourceImage || !maskImage}
                       className="flex-1"
                     >
                       {loading ? (
@@ -488,16 +509,6 @@ function InpaintingForm() {
                       {error}
                     </div>
                   )}
-
-                  {loading && jobInfo && (
-                    <div className="rounded-md bg-muted p-3 text-sm">
-                      <p>Job ID: {jobInfo.id || jobInfo.request_id || 'N/A'}</p>
-                      <p>Status: {jobInfo.status}</p>
-                      {jobInfo.progress !== undefined && (
-                        <p>Progress: {Math.round(jobInfo.progress * 100)}%</p>
-                      )}
-                    </div>
-                  )}
                 </form>
               </CardContent>
             </Card>

+ 10 - 5
webui/app/layout.tsx

@@ -4,6 +4,9 @@ import "./globals.css";
 import { ThemeProvider } from "@/components/layout";
 import { AuthProvider } from "@/lib/auth-context";
 import { ModelSelectionProvider } from "@/contexts/model-selection-context";
+import { ErrorBoundary } from "@/components/ui/error-boundary";
+import "@/lib/error-handlers";
+import "@/lib/memory-utils";
 
 const inter = Inter({
   subsets: ["latin"],
@@ -34,11 +37,13 @@ export default function RootLayout({
           enableSystem
           disableTransitionOnChange
         >
-          <AuthProvider>
-            <ModelSelectionProvider>
-              {children}
-            </ModelSelectionProvider>
-          </AuthProvider>
+          <ErrorBoundary>
+            <AuthProvider>
+              <ModelSelectionProvider>
+                {children}
+              </ModelSelectionProvider>
+            </AuthProvider>
+          </ErrorBoundary>
         </ThemeProvider>
       </body>
     </html>

+ 43 - 28
webui/app/text2img/page.tsx

@@ -1,18 +1,18 @@
 'use client';
 
 import { useState, useEffect } from 'react';
-import { Header } from '@/components/header';
+import { Header } from '@/components/layout';
 import { AppLayout } from '@/components/layout';
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Textarea } from '@/components/ui/textarea';
-import { PromptTextarea } from '@/components/prompt-textarea';
+import { PromptTextarea } from '@/components/forms';
 import { Label } from '@/components/ui/label';
 import { Card, CardContent } from '@/components/ui/card';
 import { apiClient, type GenerationRequest, type JobInfo, type ModelInfo } from '@/lib/api';
 import { Loader2, Download, X, Trash2, RotateCcw, Power } from 'lucide-react';
 import { downloadImage, downloadAuthenticatedImage } from '@/lib/utils';
-import { useLocalStorage } from '@/lib/hooks';
+import { useLocalStorage } from '@/lib/storage';
 
 const defaultFormData: GenerationRequest = {
   prompt: '',
@@ -31,7 +31,8 @@ function Text2ImgForm() {
 
   const [formData, setFormData] = useLocalStorage<GenerationRequest>(
     'text2img-form-data',
-    defaultFormData
+    defaultFormData,
+    { excludeLargeData: true, maxSize: 512 * 1024 } // 512KB limit
   );
 
   const [loading, setLoading] = useState(false);
@@ -42,6 +43,16 @@ function Text2ImgForm() {
   const [loraModels, setLoraModels] = useState<string[]>([]);
   const [embeddings, setEmbeddings] = useState<string[]>([]);
   const [error, setError] = useState<string | null>(null);
+  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
+
+  // Cleanup polling on unmount
+  useEffect(() => {
+    return () => {
+      if (pollCleanup) {
+        pollCleanup();
+      }
+    };
+  }, [pollCleanup]);
 
   useEffect(() => {
     const loadOptions = async () => {
@@ -77,10 +88,13 @@ function Text2ImgForm() {
   };
 
   const pollJobStatus = async (jobId: string) => {
-    const maxAttempts = 300; // 5 minutes with 1 second interval
+    const maxAttempts = 300; // 5 minutes with 2 second interval
     let attempts = 0;
+    let isPolling = true;
 
     const poll = async () => {
+      if (!isPolling) return;
+      
       try {
         const status = await apiClient.getJobStatus(jobId);
         setJobInfo(status);
@@ -98,42 +112,47 @@ function Text2ImgForm() {
           } else if (status.result?.images && status.result.images.length > 0) {
             // Old format: convert image URLs to authenticated URLs
             imageUrls = status.result.images.map((imageUrl: string) => {
-              // Extract filename from URL if it's already a full URL
-              if (imageUrl.includes('/output/')) {
-                const parts = imageUrl.split('/output/');
-                if (parts.length === 2) {
-                  const filename = parts[1].split('?')[0]; // Remove query params
-                  return apiClient.getImageUrl(jobId, filename);
-                }
-              }
-              // If it's just a filename, convert it directly
-              return apiClient.getImageUrl(jobId, imageUrl);
+              // Extract filename from URL if it's a full URL
+              const filename = imageUrl.split('/').pop() || imageUrl;
+              return apiClient.getImageUrl(jobId, filename);
             });
           }
 
           // Create a new array to trigger React re-render
           setGeneratedImages([...imageUrls]);
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'failed') {
           setError(status.error || 'Generation failed');
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'cancelled') {
           setError('Generation was cancelled');
           setLoading(false);
+          isPolling = false;
         } else if (attempts < maxAttempts) {
           attempts++;
           setTimeout(poll, 2000);
         } else {
           setError('Job polling timeout');
           setLoading(false);
+          isPolling = false;
         }
       } catch (err) {
-        setError(err instanceof Error ? err.message : 'Failed to check job status');
-        setLoading(false);
+        if (isPolling) {
+          setError(err instanceof Error ? err.message : 'Failed to check job status');
+          setLoading(false);
+          isPolling = false;
+        }
       }
     };
 
     poll();
+
+    // Return cleanup function
+    return () => {
+      isPolling = false;
+    };
   };
 
   const handleGenerate = async (e: React.FormEvent) => {
@@ -153,7 +172,8 @@ function Text2ImgForm() {
       setJobInfo(job);
       const jobId = job.request_id || job.id;
       if (jobId) {
-        await pollJobStatus(jobId);
+        const cleanup = pollJobStatus(jobId);
+        setPollCleanup(() => cleanup);
       } else {
         setError('No job ID returned from server');
         setLoading(false);
@@ -171,6 +191,11 @@ function Text2ImgForm() {
         await apiClient.cancelJob(jobId);
         setLoading(false);
         setError('Generation cancelled');
+        // Cleanup polling
+        if (pollCleanup) {
+          pollCleanup();
+          setPollCleanup(null);
+        }
       } catch (err) {
         console.error('Failed to cancel job:', err);
       }
@@ -423,16 +448,6 @@ function Text2ImgForm() {
                     {error}
                   </div>
                 )}
-
-                {loading && jobInfo && (
-                  <div className="rounded-md bg-muted p-3 text-sm">
-                    <p>Job ID: {jobInfo.id || jobInfo.request_id || 'N/A'}</p>
-                    <p>Status: {jobInfo.status}</p>
-                    {jobInfo.progress !== undefined && (
-                      <p>Progress: {Math.round(jobInfo.progress * 100)}%</p>
-                    )}
-                  </div>
-                )}
               </form>
             </CardContent>
           </Card>

+ 95 - 89
webui/app/upscaler/page.tsx

@@ -1,28 +1,26 @@
 'use client';
 
 import { useState, useRef, useEffect } from 'react';
-import { Header } from '@/components/header';
+import { Header } from '@/components/layout';
 import { AppLayout } from '@/components/layout';
 import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { Label } from '@/components/ui/label';
-import { Card, CardContent } from '@/components/ui/card';
+import { Card, CardContent, CardHeader, CardTitle, CardDescription } from '@/components/ui/card';
 import { apiClient, type JobInfo, type ModelInfo } from '@/lib/api';
 import { Loader2, Download, X, Upload } from 'lucide-react';
 import { downloadImage, downloadAuthenticatedImage, fileToBase64 } from '@/lib/utils';
-import { useLocalStorage } from '@/lib/hooks';
+import { useLocalStorage } from '@/lib/storage';
 import { ModelSelectionProvider, useModelSelection, useModelTypeSelection } from '@/contexts/model-selection-context';
-import { EnhancedModelSelect, EnhancedModelSelectGroup } from '@/components/enhanced-model-select';
-import { ModelSelectionWarning, AutoSelectionStatus } from '@/components/model-selection-indicator';
+import { EnhancedModelSelect } from '@/components/features/models';
+// import { AutoSelectionStatus } from '@/components/features/models';
 
 type UpscalerFormData = {
-  image: string;
   upscale_factor: number;
   model: string;
 };
 
 const defaultFormData: UpscalerFormData = {
-  image: '',
   upscale_factor: 2,
   model: '',
 };
@@ -42,16 +40,30 @@ function UpscalerForm() {
 
   const [formData, setFormData] = useLocalStorage<UpscalerFormData>(
     'upscaler-form-data',
-    defaultFormData
+    defaultFormData,
+    { excludeLargeData: true, maxSize: 512 * 1024 }
   );
 
+  // Separate state for image data (not stored in localStorage)
+  const [uploadedImage, setUploadedImage] = useState<string>('');
+  const [previewImage, setPreviewImage] = useState<string | null>(null);
+
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [jobInfo, setJobInfo] = useState<JobInfo | null>(null);
   const [generatedImages, setGeneratedImages] = useState<string[]>([]);
-  const [previewImage, setPreviewImage] = useState<string | null>(null);
+  const [pollCleanup, setPollCleanup] = useState<(() => void) | null>(null);
   const fileInputRef = useRef<HTMLInputElement>(null);
 
+  // Cleanup polling on unmount
+  useEffect(() => {
+    return () => {
+      if (pollCleanup) {
+        pollCleanup();
+      }
+    };
+  }, [pollCleanup]);
+
   useEffect(() => {
     const loadModels = async () => {
       try {
@@ -101,7 +113,7 @@ function UpscalerForm() {
 
     try {
       const base64 = await fileToBase64(file);
-      setFormData((prev) => ({ ...prev, image: base64 }));
+      setUploadedImage(base64);
       setPreviewImage(base64);
       setError(null);
     } catch (err) {
@@ -112,8 +124,11 @@ function UpscalerForm() {
   const pollJobStatus = async (jobId: string) => {
     const maxAttempts = 300;
     let attempts = 0;
+    let isPolling = true;
 
     const poll = async () => {
+      if (!isPolling) return;
+      
       try {
         const status = await apiClient.getJobStatus(jobId);
         setJobInfo(status);
@@ -147,32 +162,44 @@ function UpscalerForm() {
           // Create a new array to trigger React re-render
           setGeneratedImages([...imageUrls]);
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'failed') {
           setError(status.error || 'Upscaling failed');
           setLoading(false);
+          isPolling = false;
         } else if (status.status === 'cancelled') {
           setError('Upscaling was cancelled');
           setLoading(false);
+          isPolling = false;
         } else if (attempts < maxAttempts) {
           attempts++;
           setTimeout(poll, 2000);
         } else {
           setError('Job polling timeout');
           setLoading(false);
+          isPolling = false;
         }
       } catch (err) {
-        setError(err instanceof Error ? err.message : 'Failed to check job status');
-        setLoading(false);
+        if (isPolling) {
+          setError(err instanceof Error ? err.message : 'Failed to check job status');
+          setLoading(false);
+          isPolling = false;
+        }
       }
     };
 
     poll();
+
+    // Return cleanup function
+    return () => {
+      isPolling = false;
+    };
   };
 
   const handleUpscale = async (e: React.FormEvent) => {
     e.preventDefault();
 
-    if (!formData.image) {
+    if (!uploadedImage) {
       setError('Please upload an image first');
       return;
     }
@@ -194,12 +221,14 @@ function UpscalerForm() {
       const job = await apiClient.generateImage({
         prompt: `upscale ${formData.upscale_factor}x`,
         model: selectedUpscalerModel,
+        image: uploadedImage,
         // Add upscale-specific parameters here based on your API
       } as any);
       setJobInfo(job);
       const jobId = job.request_id || job.id;
       if (jobId) {
-        await pollJobStatus(jobId);
+        const cleanup = pollJobStatus(jobId);
+        setPollCleanup(() => cleanup);
       } else {
         setError('No job ID returned from server');
         setLoading(false);
@@ -217,6 +246,11 @@ function UpscalerForm() {
         await apiClient.cancelJob(jobId);
         setLoading(false);
         setError('Upscaling cancelled');
+        // Cleanup polling
+        if (pollCleanup) {
+          pollCleanup();
+          setPollCleanup(null);
+        }
       } catch (err) {
         console.error('Failed to cancel job:', err);
       }
@@ -237,88 +271,79 @@ function UpscalerForm() {
                   <div className="space-y-4">
                     {previewImage && (
                       <div className="relative">
-                        <img
-                          src={previewImage}
-                          alt="Source"
-                          className="w-full rounded-lg border border-border"
+                        <img 
+                          src={previewImage} 
+                          alt="Preview" 
+                          className="h-64 w-full rounded-lg object-cover"
                         />
+                        <Button
+                          type="button"
+                          variant="destructive"
+                          size="icon"
+                          className="absolute top-2 right-2 h-8 w-8"
+                          onClick={() => {
+                            setPreviewImage(null);
+                            setUploadedImage('');
+                          }}
+                        >
+                          <X className="h-4 w-4" />
+                        </Button>
                       </div>
                     )}
-                    <Button
-                      type="button"
-                      variant="outline"
-                      onClick={() => fileInputRef.current?.click()}
-                      className="w-full"
-                    >
-                      <Upload className="h-4 w-4" />
-                      {previewImage ? 'Change Image' : 'Upload Image'}
-                    </Button>
-                    <input
-                      ref={fileInputRef}
-                      type="file"
-                      accept="image/*"
-                      onChange={handleImageUpload}
-                      className="hidden"
-                    />
+                    
+                    <div className="space-y-2">
+                      <div className="flex items-center justify-center">
+                        <input
+                          ref={fileInputRef}
+                          type="file"
+                          accept="image/*"
+                          onChange={handleImageUpload}
+                          className="hidden"
+                        />
+                        <Button
+                          type="button"
+                          variant="outline"
+                          onClick={() => fileInputRef.current?.click()}
+                        >
+                          <Upload className="mr-2 h-4 w-4" />
+                          Choose Image
+                        </Button>
+                      </div>
+                    </div>
                   </div>
                 </div>
 
                 <div className="space-y-2">
-                  <Label htmlFor="upscale_factor">Upscale Factor</Label>
+                  <Label>Upscaling Factor</Label>
                   <select
-                    id="upscale_factor"
-                    name="upscale_factor"
                     value={formData.upscale_factor}
-                    onChange={handleInputChange}
-                    className="flex h-10 w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
+                    onChange={(e) => setFormData(prev => ({ ...prev, upscale_factor: Number(e.target.value) }))}
+                    className="w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
                   >
                     <option value={2}>2x (Double)</option>
                     <option value={3}>3x (Triple)</option>
                     <option value={4}>4x (Quadruple)</option>
                   </select>
-                  <p className="text-xs text-muted-foreground">
-                    Higher factors take longer to process
-                  </p>
                 </div>
 
-                {/* Model Selection Section */}
-                <EnhancedModelSelectGroup
-                  title="Model Selection"
-                  description="Select the upscaler model for image enhancement"
-                >
+                <div className="space-y-2">
+                  <Label>Upscaler Model</Label>
                   <EnhancedModelSelect
                     modelType="upscaler"
                     label="Upscaling Model"
-                    description="Model to use for upscaling the image"
-                    value={selectedUpscalerModel}
+                    value={formData.model}
+                    onValueChange={(value) => setFormData(prev => ({ ...prev, model: value }))}
+                    onSetUserOverride={(value) => {/* User override logic */}}
+                    onClearOverride={clearUpscalerUserOverride}
                     availableModels={upscalerModels}
                     isAutoSelected={isUpscalerAutoSelected}
                     isUserOverride={isUpscalerUserOverride}
-                    isLoading={state.isLoading}
-                    onValueChange={setSelectedUpscalerModel}
-                    onSetUserOverride={setUpscalerUserOverride}
-                    onClearOverride={clearUpscalerUserOverride}
                     placeholder="Select an upscaler model"
                   />
-
-                  {/* Auto-selection Status */}
-                  <div className="pt-2">
-                    <AutoSelectionStatus
-                      isAutoSelecting={state.isAutoSelecting}
-                      hasAutoSelection={Object.keys(state.autoSelectedModels).length > 0}
-                    />
-                  </div>
-
-                  {/* Warnings and Errors */}
-                  <ModelSelectionWarning
-                    warnings={state.warnings}
-                    errors={error ? [error] : []}
-                    onClearWarnings={actions.clearWarnings}
-                  />
-                </EnhancedModelSelectGroup>
+                </div>
 
                 <div className="flex gap-2">
-                  <Button type="submit" disabled={loading || !formData.image} className="flex-1">
+                  <Button type="submit" disabled={loading || !uploadedImage} className="flex-1">
                     {loading ? (
                       <>
                         <Loader2 className="h-4 w-4 animate-spin" />
@@ -341,28 +366,9 @@ function UpscalerForm() {
                     {error}
                   </div>
                 )}
-
-                {loading && jobInfo && (
-                  <div className="rounded-md bg-muted p-3 text-sm">
-                    <p>Job ID: {jobInfo.id || jobInfo.request_id || 'N/A'}</p>
-                    <p>Status: {jobInfo.status}</p>
-                    {jobInfo.progress !== undefined && (
-                      <p>Progress: {Math.round(jobInfo.progress * 100)}%</p>
-                    )}
-                  </div>
-                )}
-
-                <div className="rounded-md bg-blue-500/10 p-3 text-sm text-blue-600 dark:text-blue-400">
-                  <p className="font-medium">Note</p>
-                  <p className="mt-1">
-                    Upscaling functionality depends on your backend configuration and available upscaler models.
-                  </p>
-                </div>
               </form>
             </CardContent>
           </Card>
-
-          {/* Right Panel - Upscaled Images */}
           <Card>
             <CardContent className="pt-6">
               <div className="space-y-4">

+ 58 - 0
webui/components/ui/error-boundary.tsx

@@ -0,0 +1,58 @@
+'use client';
+
+import React, { Component, ErrorInfo, ReactNode } from 'react';
+import { AlertTriangle } from 'lucide-react';
+
+interface Props {
+  children: ReactNode;
+}
+
+interface State {
+  hasError: boolean;
+  error?: Error;
+}
+
+export class ErrorBoundary extends Component<Props, State> {
+  constructor(props: Props) {
+    super(props);
+    this.state = { hasError: false };
+  }
+
+  static getDerivedStateFromError(error: Error): State {
+    return { hasError: true, error };
+  }
+
+  componentDidCatch(error: Error, errorInfo: ErrorInfo) {
+    console.error('Error caught by boundary:', error, errorInfo);
+  }
+
+  render() {
+    if (this.state.hasError) {
+      return (
+        <div className="min-h-screen flex items-center justify-center p-4">
+          <div className="text-center max-w-md">
+            <AlertTriangle className="h-12 w-12 text-red-500 mx-auto mb-4" />
+            <h2 className="text-xl font-semibold mb-2">Something went wrong</h2>
+            <p className="text-muted-foreground mb-4">
+              The application encountered an unexpected error. Try refreshing the page.
+            </p>
+            {process.env.NODE_ENV === 'development' && this.state.error && (
+              <details className="text-left text-sm text-red-600 bg-red-50 p-4 rounded-lg mb-4">
+                <summary className="cursor-pointer font-medium mb-2">Error details</summary>
+                <pre className="whitespace-pre-wrap">{this.state.error.stack}</pre>
+              </details>
+            )}
+            <button
+              onClick={() => window.location.reload()}
+              className="px-4 py-2 bg-primary text-primary-foreground rounded-lg hover:bg-primary/90 transition-colors"
+            >
+              Refresh Page
+            </button>
+          </div>
+        </div>
+      );
+    }
+
+    return this.props.children;
+  }
+}

+ 69 - 0
webui/lib/error-handlers.ts

@@ -0,0 +1,69 @@
+// Global error handlers to prevent freezing
+if (typeof window !== 'undefined') {
+  // Handle unhandled promise rejections
+  window.addEventListener('unhandledrejection', (event) => {
+    console.error('Unhandled promise rejection:', event.reason);
+    
+    // Prevent the error from causing the app to freeze
+    event.preventDefault();
+    
+    // Optionally show a user-friendly notification
+    if (event.reason instanceof Error && event.reason.message) {
+      console.warn('Application error:', event.reason.message);
+    }
+  });
+
+  // Handle uncaught errors
+  window.addEventListener('error', (event) => {
+    console.error('Uncaught error:', event.error);
+    
+    // Prevent the error from causing the app to freeze
+    event.preventDefault();
+  });
+
+  // Handle resource loading errors
+  window.addEventListener('error', (event) => {
+    if (event.target !== window) {
+      console.warn('Resource loading error:', event.target);
+      event.preventDefault();
+    }
+  }, true);
+}
+
+// Export a utility function for manual error reporting
+export function reportError(error: Error, context?: string) {
+  console.error(`Application error${context ? ` in ${context}` : ''}:`, error);
+  
+  // In development, you might want to show more details
+  if (process.env.NODE_ENV === 'development') {
+    console.trace('Error stack trace');
+  }
+}
+
+// Export a function to safely run async operations
+export async function safeAsync<T>(
+  operation: () => Promise<T>,
+  fallback?: T,
+  context?: string
+): Promise<T | undefined> {
+  try {
+    return await operation();
+  } catch (error) {
+    reportError(error instanceof Error ? error : new Error(String(error)), context);
+    return fallback;
+  }
+}
+
+// Export a function to safely run sync operations
+export function safeSync<T>(
+  operation: () => T,
+  fallback?: T,
+  context?: string
+): T | undefined {
+  try {
+    return operation();
+  } catch (error) {
+    reportError(error instanceof Error ? error : new Error(String(error)), context);
+    return fallback;
+  }
+}

+ 125 - 0
webui/lib/memory-utils.ts

@@ -0,0 +1,125 @@
+// Memory leak prevention utilities
+
+class RequestTracker {
+  private activeRequests = new Set<AbortController>();
+
+  addRequest(controller: AbortController) {
+    this.activeRequests.add(controller);
+  }
+
+  removeRequest(controller: AbortController) {
+    this.activeRequests.delete(controller);
+  }
+
+  cancelAllRequests() {
+    this.activeRequests.forEach(controller => {
+      try {
+        controller.abort();
+      } catch (error) {
+        console.warn('Failed to abort request:', error);
+      }
+    });
+    this.activeRequests.clear();
+  }
+
+  getActiveRequestCount(): number {
+    return this.activeRequests.size;
+  }
+}
+
+export const requestTracker = new RequestTracker();
+
+// Enhanced fetch with automatic tracking
+export function trackedFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
+  const controller = new AbortController();
+  const timeoutId = setTimeout(() => {
+    controller.abort();
+  }, 30000); // 30 second timeout
+
+  // Merge abort signals
+  if (init?.signal) {
+    init.signal.addEventListener('abort', () => {
+      controller.abort();
+    });
+  }
+
+  const trackedInit = {
+    ...init,
+    signal: controller.signal,
+  };
+
+  requestTracker.addRequest(controller);
+
+  return fetch(input, trackedInit).finally(() => {
+    clearTimeout(timeoutId);
+    requestTracker.removeRequest(controller);
+  });
+}
+
+// Cleanup utility for components
+export function useCleanup() {
+  const cleanupFunctions = new Set<() => void>();
+
+  const addCleanup = (cleanup: () => void) => {
+    cleanupFunctions.add(cleanup);
+  };
+
+  const executeCleanup = () => {
+    cleanupFunctions.forEach(cleanup => {
+      try {
+        cleanup();
+      } catch (error) {
+        console.warn('Cleanup function failed:', error);
+      }
+    });
+    cleanupFunctions.clear();
+  };
+
+  return { addCleanup, executeCleanup };
+}
+
+// Memory monitoring for development
+export function startMemoryMonitoring() {
+  if (process.env.NODE_ENV === 'development' && typeof window !== 'undefined') {
+    setInterval(() => {
+      if ('memory' in performance) {
+        const memory = (performance as any).memory;
+        const usedMB = Math.round(memory.usedJSHeapSize / 1048576);
+        const totalMB = Math.round(memory.totalJSHeapSize / 1048576);
+        const limitMB = Math.round(memory.jsHeapSizeLimit / 1048576);
+        
+        // Log memory usage every 30 seconds
+        if (usedMB > limitMB * 0.9) {
+          console.warn(`High memory usage: ${usedMB}MB / ${limitMB}MB`);
+        }
+      }
+    }, 30000);
+  }
+}
+
+// Debounced function with cleanup
+export function debounceWithCleanup<T extends (...args: any[]) => any>(
+  func: T,
+  wait: number
+): { debouncedFn: (...args: Parameters<T>) => void; cancel: () => void } {
+  let timeoutId: NodeJS.Timeout | null = null;
+
+  const debouncedFn = (...args: Parameters<T>) => {
+    if (timeoutId) {
+      clearTimeout(timeoutId);
+    }
+    timeoutId = setTimeout(() => {
+      func(...args);
+      timeoutId = null;
+    }, wait);
+  };
+
+  const cancel = () => {
+    if (timeoutId) {
+      clearTimeout(timeoutId);
+      timeoutId = null;
+    }
+  };
+
+  return { debouncedFn, cancel };
+}

+ 136 - 0
webui/lib/storage.ts

@@ -0,0 +1,136 @@
+import { useState, useEffect } from 'react';
+
+// Safe localStorage hook that handles quota errors
+export function useLocalStorage<T>(
+  key: string,
+  initialValue: T,
+  options?: {
+    excludeLargeData?: boolean;
+    maxSize?: number; // in bytes
+  }
+) {
+  const { excludeLargeData = true, maxSize = 1024 * 1024 } = options || {}; // 1MB default
+
+  // Get stored value
+  const getStoredValue = (): T => {
+    if (typeof window === 'undefined') {
+      return initialValue;
+    }
+
+    try {
+      const item = window.localStorage.getItem(key);
+      if (item === null) {
+        return initialValue;
+      }
+      return JSON.parse(item);
+    } catch (error) {
+      console.warn(`Error reading localStorage key "${key}":`, error);
+      
+      // If there's an error, clear the key and return initial value
+      try {
+        window.localStorage.removeItem(key);
+      } catch (clearError) {
+        console.warn(`Error clearing localStorage key "${key}":`, clearError);
+      }
+      
+      return initialValue;
+    }
+  };
+
+  const [storedValue, setStoredValue] = useState<T>(getStoredValue);
+
+  // Set value to localStorage
+  const setValue = (value: T | ((val: T) => T)) => {
+    try {
+      const valueToStore = value instanceof Function ? value(storedValue) : value;
+      
+      // Check if we should exclude large data
+      if (excludeLargeData) {
+        const serialized = JSON.stringify(valueToStore);
+        
+        // Check for large base64 images or data URLs
+        if (serialized.includes('data:image/') || serialized.length > maxSize) {
+          console.warn(`Skipping localStorage save for "${key}" - data too large or contains images`);
+          setStoredValue(valueToStore);
+          return;
+        }
+      }
+
+      setStoredValue(valueToStore);
+      
+      if (typeof window !== 'undefined') {
+        window.localStorage.setItem(key, JSON.stringify(valueToStore));
+      }
+    } catch (error) {
+      console.warn(`Error saving localStorage key "${key}":`, error);
+      
+      // Still update the state even if localStorage fails
+      const valueToStore = value instanceof Function ? value(storedValue) : value;
+      setStoredValue(valueToStore);
+      
+      // Try to clear some space if it's a quota error
+      if (error instanceof Error && error.name === 'QuotaExceededError') {
+        try {
+          // Clear some non-essential keys
+          const keysToClear = ['img2img-form-data', 'text2img-form-data', 'inpainting-form-data'];
+          keysToClear.forEach(k => {
+            if (k !== key) {
+              window.localStorage.removeItem(k);
+            }
+          });
+        } catch (clearError) {
+          console.warn('Error clearing localStorage space:', clearError);
+        }
+      }
+    }
+  };
+
+  return [storedValue, setValue] as const;
+}
+
+// Session storage alternative for large data
+export function useSessionStorage<T>(key: string, initialValue: T) {
+  const [storedValue, setStoredValue] = useState<T>(() => {
+    if (typeof window === 'undefined') {
+      return initialValue;
+    }
+
+    try {
+      const item = window.sessionStorage.getItem(key);
+      return item ? JSON.parse(item) : initialValue;
+    } catch (error) {
+      console.warn(`Error reading sessionStorage key "${key}":`, error);
+      return initialValue;
+    }
+  });
+
+  const setValue = (value: T | ((val: T) => T)) => {
+    try {
+      const valueToStore = value instanceof Function ? value(storedValue) : value;
+      setStoredValue(valueToStore);
+      
+      if (typeof window !== 'undefined') {
+        window.sessionStorage.setItem(key, JSON.stringify(valueToStore));
+      }
+    } catch (error) {
+      console.warn(`Error saving sessionStorage key "${key}":`, error);
+      // Still update the state even if sessionStorage fails
+      const valueToStore = value instanceof Function ? value(storedValue) : value;
+      setStoredValue(valueToStore);
+    }
+  };
+
+  return [storedValue, setValue] as const;
+}
+
+// Memory-only storage for large data
+export function useMemoryStorage<T>(initialValue: T) {
+  const [storedValue, setStoredValue] = useState<T>(initialValue);
+
+  const setValue = (value: T | ((val: T) => T)) => {
+    const valueToStore = value instanceof Function ? value(storedValue) : value;
+    setStoredValue(valueToStore);
+  };
+
+  return [storedValue, setValue] as const;
+}