着手增加逻辑

app.py

@@ -1,11 +1,20 @@
-import spaces
+import os
+import gc
 import gradio as gr
 import numpy as np
 import torch
+import json
+import spaces
 import random
-import logging
+import config
 import utils
+import logging
+from PIL import Image, PngImagePlugin
+from datetime import datetime
 from diffusers.models import AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
+import time
+from typing import List, Dict, Tuple, Optional
 from config import (
     MODEL,
     MIN_IMAGE_SIZE,
@@ -52,6 +61,37 @@ else:
     pipe = None
 
 
+class GenerationError(Exception):
+    """Custom exception for generation errors"""
+    pass
+
+def validate_prompt(prompt: str) -> str:
+    """Validate and clean up the input prompt."""
+    if not isinstance(prompt, str):
+        raise GenerationError("Prompt must be a string")
+    try:
+        # Ensure proper UTF-8 encoding/decoding
+        prompt = prompt.encode('utf-8').decode('utf-8')
+        # Add space between ! and ,
+        prompt = prompt.replace("!,", "! ,")
+    except UnicodeError:
+        raise GenerationError("Invalid characters in prompt")
+    
+    # Only check if the prompt is completely empty or only whitespace
+    if not prompt or prompt.isspace():
+        raise GenerationError("Prompt cannot be empty")
+    return prompt.strip()
+
+def validate_dimensions(width: int, height: int) -> None:
+    """Validate image dimensions."""
+    if not MIN_IMAGE_SIZE <= width <= MAX_IMAGE_SIZE:
+        raise GenerationError(f"Width must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
+        
+    if not MIN_IMAGE_SIZE <= height <= MAX_IMAGE_SIZE:
+        raise GenerationError(f"Height must be between {MIN_IMAGE_SIZE} and {MAX_IMAGE_SIZE}")
+
+
+
 
 @spaces.GPU
 def generate(
@@ -70,23 +110,39 @@ def generate(
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    
+    """Generate images based on the given parameters."""
+    start_time = time.time()
+    upscaler_pipe = None
+    backup_scheduler = None
 
-    # generator = torch.Generator().manual_seed(seed)
-
-    # image = pipe(
-    #     prompt=prompt,
-    #     negative_prompt=negative_prompt,
-    #     guidance_scale=guidance_scale,
-    #     num_inference_steps=num_inference_steps,
-    #     width=width,
-    #     height=height,
-    #     generator=generator,
-    # ).images[0]
+    try:
+        # Memory management
+        torch.cuda.empty_cache()
+        gc.collect()
 
-    # return image, seed
+        return None
+    except GenerationError as e:
+        logger.warning(f"Generation validation error: {str(e)}")
+        raise gr.Error(str(e))
+    except Exception as e:
+        logger.exception("Unexpected error during generation")
+        raise gr.Error(f"Generation failed: {str(e)}")
+    finally:
+        # Cleanup
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        if upscaler_pipe is not None:
+            del upscaler_pipe
+        
+        if backup_scheduler is not None and pipe is not None:
+            pipe.scheduler = backup_scheduler
+            
+        utils.free_memory()
 
-    return None, seed
 
+   
 
 
 
@@ -207,7 +263,7 @@ with gr.Blocks(css=custom_css).queue() as demo:
             seed,randomize_seed,
             guidance_scale,num_inference_steps
             ], 
-        outputs=[result, seed],
+        outputs=[result],
     )  
 
 if __name__ == "__main__":