着手加入模型

app.py

@@ -28,21 +28,21 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
 # Model initialization
-# if torch.cuda.is_available():
-#     try:
-#         logger.info("Loading VAE and pipeline...")
-#         vae = AutoencoderKL.from_pretrained(
-#             "madebyollin/sdxl-vae-fp16-fix",
-#             torch_dtype=torch.float16,
-#         )
-#         pipe = utils.load_pipeline("cagliostrolab/animagine-xl-4.0", device, vae=vae)
-#         logger.info("Pipeline loaded successfully on GPU!")
-#     except Exception as e:
-#         logger.error(f"Error loading VAE, falling back to default: {e}")
-#         pipe = utils.load_pipeline("cagliostrolab/animagine-xl-4.0", device)
-# else:
-#     logger.warning("CUDA not available, running on CPU")
-#     pipe = None
+if torch.cuda.is_available():
+    try:
+        logger.info("Loading VAE and pipeline...")
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=torch.float16,
+        )
+        pipe = utils.load_pipeline("cagliostrolab/animagine-xl-4.0", device, vae=vae)
+        logger.info("Pipeline loaded successfully on GPU!")
+    except Exception as e:
+        logger.error(f"Error loading VAE, falling back to default: {e}")
+        pipe = utils.load_pipeline("cagliostrolab/animagine-xl-4.0", device)
+else:
+    logger.warning("CUDA not available, running on CPU")
+    pipe = None
 
 
 
@@ -211,4 +211,5 @@ with gr.Blocks(css=custom_css).queue() as demo:
         outputs=[result, seed],
     )  
 
-demo.launch()
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

utils.py

@@ -1,5 +1,14 @@
+import gc
+import os
+import random
+import numpy as np
+import json
 import torch
-from typing import  Optional, Any
+import uuid
+from PIL import Image, PngImagePlugin
+from datetime import datetime
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple, Any, List
 from diffusers import (
     DDIMScheduler,
     DPMSolverMultistepScheduler,