Update app.py

app.py

@@ -44,32 +44,47 @@ def load_model(model_id):
         logger.info(f"Loading model {model_id}...")
         try:
             if model_id == "ssd-1b":
-                # 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"
+                # Try loading UNet from xet/unet
+                try:
+                    logger.info(f"Preloading UNet for {model_id} from {model_paths[model_id]}/xet/unet")
+                    unet_config = UNet2DConditionModel.load_config(
+                        f"{model_paths[model_id]}/xet/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]}/xet/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,
+                        force_download=True
+                    )
+                except Exception as e:
+                    logger.warning(f"Failed to load UNet for {model_id}: {str(e)}")
+                    logger.info(f"Falling back to standard pipeline loading for {model_id}")
+                    # Fallback to standard pipeline
+                    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
                     )
-                )
-                # 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(
@@ -77,7 +92,8 @@ def load_model(model_id):
                     torch_dtype=torch.float32,
                     use_auth_token=os.getenv("HF_TOKEN"),
                     use_safetensors=True,
-                    low_cpu_mem_usage=True
+                    low_cpu_mem_usage=True,
+                    force_download=True
                 )
             logger.info(f"Pipeline components loading for {model_id}...")
             pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
@@ -123,7 +139,7 @@ def generate():
         pipe.to(torch.device("cpu"))
 
         images = []
-        num_inference_steps = 10 if model_id == 'ssd-1b' else 30
+        num_inference_steps = 20 if model_id == 'ssd-1b' else 30
         for _ in range(num_images):
             image = pipe(
                 prompt=prompt,