Update app.py

app.py

@@ -1,6 +1,6 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from diffusers import StableDiffusionPipeline, DiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 import torch
 import os
 from PIL import Image
@@ -9,106 +9,114 @@ import time
 from accelerate import Accelerator
 import logging
 
-app = Flask(__name__)
+# Disable GPU detection
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+torch.set_default_device("cpu")
+
+app = Flask(__name__, static_folder='static')
 CORS(app)
 
-     # Configure logging
+# Configure logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-     # Initialize Accelerator for CPU
-accelerator = Accelerator(cpu=True)
+# Initialize Accelerator for CPU
+accelerator = Accelerator(device_placement=False)
 
-     # Model cache
+# Model cache
 model_cache = {}
 model_paths = {
-         "ssd-1b": "remiai3/ssd-1b",
-         "sd-v1-5": "remiai3/stable-diffusion-v1-5"
-     }
+    "ssd-1b": "remiai3/ssd-1b",
+    "sd-v1-5": "remiai3/stable-diffusion-v1-5"
+}
 
-     # Image ratio to dimensions (optimized for CPU)
+# Image ratio to dimensions (optimized for CPU)
 ratio_to_dims = {
-         "1:1": (256, 256),
-         "3:4": (192, 256),
-         "16:9": (256, 144)
-     }
+    "1:1": (256, 256),
+    "3:4": (192, 256),
+    "16:9": (256, 144)
+}
 
 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
-                 )
-                 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-                 pipe = accelerator.prepare(pipe)
-                 pipe.enable_attention_slicing()
-                 pipe.enable_sequential_cpu_offload()
-                 model_cache[model_id] = pipe
-                 logger.info(f"Model {model_id} loaded successfully")
-             except Exception as e:
-                 logger.error(f"Error loading model {model_id}: {str(e)}")
-                 raise
-         return model_cache[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,
+                device_map="cpu"
+            )
+            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+            pipe = accelerator.prepare(pipe)
+            pipe.enable_attention_slicing()
+            pipe.enable_sequential_cpu_offload()
+            pipe.to("cpu")
+            model_cache[model_id] = pipe
+            logger.info(f"Model {model_id} loaded successfully")
+        except Exception as e:
+            logger.error(f"Error loading model {model_id}: {str(e)}")
+            raise
+    return model_cache[model_id]
 
 @app.route('/')
 def index():
-         return app.send_static_file('index.html')
+    return app.send_static_file('index.html')
 
 @app.route('/generate', methods=['POST'])
 def generate():
-         try:
-             data = request.json
-             model_id = data.get('model', 'ssd-1b')
-             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))
-
-             if not prompt:
-                 return jsonify({"error": "Prompt is required"}), 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':
-                 return jsonify({"error": "SSD-1B supports only 1:1 ratio"}), 400
-             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))
-             pipe = load_model(model_id)
-
-             images = []
-             for _ in range(num_images):
-                 image = pipe(
-                     prompt=prompt,
-                     height=height,
-                     width=width,
-                     num_inference_steps=20,
-                     guidance_scale=guidance_scale
-                 ).images[0]
-                 images.append(image)
-
-             output_dir = "outputs"
-             os.makedirs(output_dir, exist_ok=True)
-             image_urls = []
-             for i, img in enumerate(images):
-                 img_path = os.path.join(output_dir, f"generated_{int(time.time())}_{i}.png")
-                 img.save(img_path)
-                 with open(img_path, "rb") as f:
-                     img_data = base64.b64encode(f.read()).decode('utf-8')
-                 image_urls.append(f"data:image/png;base64,{img_data}")
-                 os.remove(img_path)
-
-             return jsonify({"images": image_urls})
-
-         except Exception as e:
-             logger.error(f"Image generation failed: {str(e)}")
-             return jsonify({"error": f"Image generation failed: {str(e)}"}), 500
+    try:
+        data = request.json
+        model_id = data.get('model', 'ssd-1b')
+        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))
+
+        if not prompt:
+            return jsonify({"error": "Prompt is required"}), 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':
+            return jsonify({"error": "SSD-1B supports only 1:1 ratio"}), 400
+        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))
+        pipe = load_model(model_id)
+        pipe.to("cpu")
+
+        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)
+
+        output_dir = "outputs"
+        os.makedirs(output_dir, exist_ok=True)
+        image_urls = []
+        for i, img in enumerate(images):
+            img_path = os.path.join(output_dir, f"generated_{int(time.time())}_{i}.png")
+            img.save(img_path)
+            with open(img_path, "rb") as f:
+                img_data = base64.b64encode(f.read()).decode('utf-8')
+            image_urls.append(f"data:image/png;base64,{img_data}")
+            os.remove(img_path)
+
+        return jsonify({"images": image_urls})
+
+    except Exception as e:
+        logger.error(f"Image generation failed: {str(e)}")
+        return jsonify({"error": f"Image generation failed: {str(e)}"}), 500
 
 if __name__ == '__main__':
-         app.run(host='0.0.0.0', port=7860)
+    app.run(host='0.0.0.0', port=7860)