Update app.py

app.py

@@ -1,6 +1,7 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers.models import UNet2DConditionModel
 import torch
 import os
 from PIL import Image
@@ -42,20 +43,43 @@ def load_model(model_id):
     if model_id not in model_cache:
         logger.info(f"Loading model {model_id}...")
         try:
-            pipe = StableDiffusionPipeline.from_pretrained(
-                model_paths[model_id],
-                torch_dtype=torch.float32,
-                use_auth_token=os.getenv("HF_TOKEN"),
-                use_safetensors=True,
-                low_cpu_mem_usage=True
-            )
-            logger.info(f"Pipeline components loading for {model_id}...")
             if model_id == "ssd-1b":
-                # Patch UNet config to remove problematic parameter
-                if hasattr(pipe.unet.config, "reverse_transformer_layers_per_block"):
-                    logger.info(f"Original UNet config for {model_id}: {pipe.unet.config}")
-                    pipe.unet.config.reverse_transformer_layers_per_block = None
-                    logger.info(f"Patched UNet config for {model_id}: {pipe.unet.config}")
+                # Preload UNet and patch configuration
+                logger.info(f"Preloading UNet for {model_id}")
+                unet_config = UNet2DConditionModel.load_config(
+                    f"{model_paths[model_id]}/unet",
+                    use_auth_token=os.getenv("HF_TOKEN")
+                )
+                if "reverse_transformer_layers_per_block" in unet_config:
+                    logger.info(f"Original UNet config for {model_id}: {unet_config}")
+                    unet_config["reverse_transformer_layers_per_block"] = None
+                    logger.info(f"Patched UNet config for {model_id}: {unet_config}")
+                unet = UNet2DConditionModel.from_config(unet_config)
+                unet.load_state_dict(
+                    torch.load(
+                        f"{model_paths[model_id]}/unet/diffusion_pytorch_model.bin",
+                        map_location="cpu"
+                    )
+                )
+                # Load pipeline with patched UNet
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    model_paths[model_id],
+                    unet=unet,
+                    torch_dtype=torch.float32,
+                    use_auth_token=os.getenv("HF_TOKEN"),
+                    use_safetensors=True,
+                    low_cpu_mem_usage=True
+                )
+            else:
+                # Standard loading for sd-v1-5
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    model_paths[model_id],
+                    torch_dtype=torch.float32,
+                    use_auth_token=os.getenv("HF_TOKEN"),
+                    use_safetensors=True,
+                    low_cpu_mem_usage=True
+                )
+            logger.info(f"Pipeline components loading for {model_id}...")
             pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
             pipe.enable_attention_slicing()
             pipe.to(torch.device("cpu"))
@@ -99,7 +123,7 @@ def generate():
         pipe.to(torch.device("cpu"))
 
         images = []
-        num_inference_steps = 20 if model_id == 'ssd-1b' else 30
+        num_inference_steps = 10 if model_id == 'ssd-1b' else 30
         for _ in range(num_images):
             image = pipe(
                 prompt=prompt,