Update app.py

app.py

@@ -1,15 +1,15 @@
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
-from diffusers.models import UNet2DConditionModel
 import torch
 import os
 from PIL import Image
 import base64
 import time
 import logging
+from huggingface_hub import list_repo_files
 
-# Disable GPU detection
+# Disable GPU detection (remove these lines if GPU is available)
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["CUDA_DEVICE_ORDER"] = ""
 os.environ["TORCH_CUDA_ARCH_LIST"] = ""
@@ -28,7 +28,7 @@ logger.info(f"Device in use: {torch.device('cpu')}")
 # Model cache
 model_cache = {}
 model_paths = {
-    "ssd-1b": "remiai3/ssd-1b",
+    "ssd-1b": "segmind/SSD-1B",  # Use segmind/SSD-1B for testing
     "sd-v1-5": "remiai3/stable-diffusion-v1-5"
 }
 
@@ -41,63 +41,25 @@ ratio_to_dims = {
 
 def load_model(model_id):
     if model_id not in model_cache:
-        logger.info(f"Loading model {model_id}...")
+        logger.info(f"Loading model {model_id} from {model_paths[model_id]}")
+        logger.info(f"HF_TOKEN present: {os.getenv('HF_TOKEN') is not None}")
         try:
-            if model_id == "ssd-1b":
-                # Try StableDiffusionXLPipeline first
-                try:
-                    logger.info(f"Attempting StableDiffusionXLPipeline for {model_id}")
-                    pipe = StableDiffusionXLPipeline.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,
-                        force_download=True
-                    )
-                except Exception as e:
-                    logger.warning(f"StableDiffusionXLPipeline failed for {model_id}: {str(e)}")
-                    logger.info(f"Falling back to StableDiffusionPipeline for {model_id}")
-                    # Fallback to StableDiffusionPipeline with patched UNet
-                    unet_config = UNet2DConditionModel.load_config(
-                        f"{model_paths[model_id]}/unet",
-                        use_auth_token=os.getenv("HF_TOKEN"),
-                        force_download=True
-                    )
-                    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"
-                        )
-                    )
-                    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,
-                        force_download=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,
-                    force_download=True
-                )
-            logger.info(f"Pipeline components loading for {model_id}...")
+            # Log repository files for debugging
+            repo_files = list_repo_files(model_paths[model_id], token=os.getenv("HF_TOKEN"))
+            logger.info(f"Files in {model_paths[model_id]}: {repo_files}")
+
+            # Choose pipeline based on model
+            pipe_class = StableDiffusionXLPipeline if model_id == "ssd-1b" else StableDiffusionPipeline
+            pipe = pipe_class.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.enable_attention_slicing()
-            pipe.to(torch.device("cpu"))
+            pipe.to(torch.device("cpu"))  # Change to "cuda" if GPU is available
             model_cache[model_id] = pipe
             logger.info(f"Model {model_id} loaded successfully")
         except Exception as e:
@@ -135,7 +97,7 @@ def generate():
 
         width, height = ratio_to_dims.get(ratio, (256, 256))
         pipe = load_model(model_id)
-        pipe.to(torch.device("cpu"))
+        pipe.to(torch.device("cpu"))  # Change to "cuda" if GPU is available
 
         images = []
         num_inference_steps = 20 if model_id == 'ssd-1b' else 30