Update app.py

app.py

@@ -1,6 +1,6 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, UniPCMultistepScheduler
 import torch
 import os
 from PIL import Image
@@ -8,6 +8,7 @@ import base64
 import time
 import logging
 from huggingface_hub import list_repo_files
+import psutil
 
 # Disable GPU detection (remove these lines if GPU is available)
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
@@ -22,21 +23,22 @@ CORS(app)
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Log device in use
+# Log device and memory info
 logger.info(f"Device in use: {torch.device('cpu')}")
+logger.info(f"Available memory: {psutil.virtual_memory().available / (1024 ** 3):.2f} GB")
 
 # Model cache
 model_cache = {}
 model_paths = {
-    "ssd-1b": "segmind/SSD-1B",  # Use segmind/SSD-1B for testing
+    "ssd-1b": "segmind/SSD-1B",  # Using segmind/SSD-1B for testing
     "sd-v1-5": "remiai3/stable-diffusion-v1-5"
 }
 
-# Image ratio to dimensions (optimized for CPU)
+# Image ratio to dimensions (optimized for CPU, multiple of 8)
 ratio_to_dims = {
-    "1:1": (256, 256),
-    "3:4": (192, 256),
-    "16:9": (256, 144)
+    "1:1": (512, 512),  # Default for SSD-1B
+    "3:4": (384, 512),
+    "16:9": (512, 288)
 }
 
 def load_model(model_id):
@@ -57,7 +59,9 @@ def load_model(model_id):
                 use_safetensors=True,
                 low_cpu_mem_usage=True
             )
-            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+            # Use UniPCMultistepScheduler for SSD-1B, DPMSolver for SD-v1-5
+            scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) if model_id == "ssd-1b" else DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+            pipe.scheduler = scheduler
             pipe.enable_attention_slicing()
             pipe.to(torch.device("cpu"))  # Change to "cuda" if GPU is available
             model_cache[model_id] = pipe
@@ -83,11 +87,15 @@ def generate():
         prompt = data.get('prompt', '')
         ratio = data.get('ratio', '1:1')
         num_images = min(int(data.get('num_images', 1)), 4)
-        guidance_scale = float(data.get('guidance_scale', 7.5))
+        guidance_scale = float(min(max(data.get('guidance_scale', 7.5), 1.0), 20.0))  # Clamp guidance scale
+
+        # Log input parameters
+        logger.info(f"Generating with model: {model_id}, prompt: {prompt}, ratio: {ratio}, num_images: {num_images}, guidance_scale: {guidance_scale}")
 
         if not prompt:
             return jsonify({"error": "Prompt is required"}), 400
-
+        if len(prompt) > 512:
+            return jsonify({"error": "Prompt is too long (max 512 characters)"}), 400
         if model_id == 'ssd-1b' and num_images > 1:
             return jsonify({"error": "SSD-1B allows only 1 image per generation"}), 400
         if model_id == 'ssd-1b' and ratio != '1:1':
@@ -95,21 +103,34 @@ def generate():
         if model_id == 'sd-v1-5' and len(prompt.split()) > 77:
             return jsonify({"error": "Prompt exceeds 77 tokens for Stable Diffusion v1.5"}), 400
 
-        width, height = ratio_to_dims.get(ratio, (256, 256))
+        width, height = ratio_to_dims.get(ratio, (512, 512))
+        if width % 8 != 0 or height % 8 != 0:
+            return jsonify({"error": "Width and height must be multiples of 8"}), 400
+
+        # Log memory before generation
+        logger.info(f"Memory before generation: {psutil.virtual_memory().available / (1024 ** 3):.2f} GB")
+
         pipe = load_model(model_id)
         pipe.to(torch.device("cpu"))  # Change to "cuda" if GPU is available
 
         images = []
-        num_inference_steps = 20 if model_id == 'ssd-1b' else 30
-        for _ in range(num_images):
-            image = pipe(
-                prompt=prompt,
-                height=height,
-                width=width,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale
-            ).images[0]
-            images.append(image)
+        num_inference_steps = 30 if model_id == 'ssd-1b' else 40  # Unified steps for stability
+        try:
+            for _ in range(num_images):
+                image = pipe(
+                    prompt=prompt,
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale
+                ).images[0]
+                images.append(image)
+        except IndexError as e:
+            logger.error(f"IndexError during generation: {str(e)}")
+            return jsonify({"error": f"Generation failed due to invalid index access: {str(e)}"}), 500
+        except Exception as e:
+            logger.error(f"Unexpected error during generation: {str(e)}")
+            return jsonify({"error": f"Generation failed: {str(e)}"}), 500
 
         output_dir = "outputs"
         os.makedirs(output_dir, exist_ok=True)
@@ -122,6 +143,7 @@ def generate():
             image_urls.append(f"data:image/png;base64,{img_data}")
             os.remove(img_path)
 
+        logger.info(f"Generation successful, returning {len(image_urls)} images")
         return jsonify({"images": image_urls})
 
     except Exception as e: