Update model/pipeline.py

model/pipeline.py

@@ -42,72 +42,154 @@ class CatVTONPipeline:
         self.weight_dtype = weight_dtype
         self.skip_safety_check = skip_safety_check
 
+        print(f"[INFO] Pipeline başlatılıyor - Device: {self.device}, dtype: {self.weight_dtype}")
+
         # TF32 sadece CUDA'da anlamlı
         if use_tf32 and self.device.type == "cuda":
             torch.set_float32_matmul_precision("high")
             torch.backends.cuda.matmul.allow_tf32 = True
+            print("[INFO] TF32 etkinleştirildi")
 
         # Scheduler
-        self.noise_scheduler = DDIMScheduler.from_pretrained(base_ckpt, subfolder="scheduler")
+        try:
+            self.noise_scheduler = DDIMScheduler.from_pretrained(base_ckpt, subfolder="scheduler")
+            print("[INFO] Scheduler yüklendi")
+        except Exception as e:
+            print(f"[ERROR] Scheduler yüklenirken hata: {e}")
+            raise
 
         # VAE yükleme (hata durumunda fallback)
         try:
             self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(
                 self.device, dtype=self.weight_dtype
             )
+            print("[INFO] VAE yüklendi")
         except Exception as e:
             print(f"[WARN] VAE yüklenirken hata: {e}. float32 fallback yapılıyor.")
-            self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(
-                self.device, dtype=torch.float32
-            )
+            try:
+                self.vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(
+                    self.device, dtype=torch.float32
+                )
+                print("[INFO] VAE float32 ile yüklendi")
+            except Exception as e2:
+                print(f"[ERROR] VAE yüklenemedi: {e2}")
+                raise
 
-        # Safety checker
+        # Safety checker - ZeroGPU uyumlu hata yakalama
         if not skip_safety_check:
-            self.feature_extractor = CLIPImageProcessor.from_pretrained(base_ckpt, subfolder="feature_extractor")
-            self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-                base_ckpt, subfolder="safety_checker"
-            ).to(self.device, dtype=self.weight_dtype)
+            try:
+                self.feature_extractor = CLIPImageProcessor.from_pretrained(base_ckpt, subfolder="feature_extractor")
+                
+                # Safety checker için çoklu deneme stratejisi
+                safety_checker_loaded = False
+                
+                # 1. Önce safetensors formatını dene
+                try:
+                    self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+                        base_ckpt, 
+                        subfolder="safety_checker",
+                        use_safetensors=True,
+                        torch_dtype=self.weight_dtype
+                    ).to(self.device, dtype=self.weight_dtype)
+                    safety_checker_loaded = True
+                    print("[INFO] Safety checker (safetensors) yüklendi")
+                except Exception as e1:
+                    print(f"[WARN] Safetensors safety checker yüklenemedi: {e1}")
+                
+                # 2. Eğer safetensors başarısızsa, farklı model dene
+                if not safety_checker_loaded:
+                    try:
+                        self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+                            "CompVis/stable-diffusion-safety-checker",
+                            torch_dtype=self.weight_dtype
+                        ).to(self.device, dtype=self.weight_dtype)
+                        safety_checker_loaded = True
+                        print("[INFO] Safety checker (CompVis) yüklendi")
+                    except Exception as e2:
+                        print(f"[WARN] CompVis safety checker yüklenemedi: {e2}")
+                
+                # 3. Son çare olarak safety checker'ı devre dışı bırak
+                if not safety_checker_loaded:
+                    print("[WARN] Safety checker yüklenemedi, devre dışı bırakılıyor")
+                    self.safety_checker = None
+                    self.feature_extractor = None
+                    self.skip_safety_check = True
+                    
+            except Exception as e:
+                print(f"[WARN] Safety checker başlatılamadı: {e}. Devre dışı bırakılıyor.")
+                self.feature_extractor = None
+                self.safety_checker = None
+                self.skip_safety_check = True
         else:
             self.feature_extractor = None
             self.safety_checker = None
+            print("[INFO] Safety checker devre dışı")
 
         # UNet ve adapter
-        self.unet = UNet2DConditionModel.from_pretrained(base_ckpt, subfolder="unet").to(
-            self.device, dtype=self.weight_dtype
-        )
-        init_adapter(self.unet, cross_attn_cls=SkipAttnProcessor)
-        self.attn_modules = get_trainable_module(self.unet, "attention")
-        self.auto_attn_ckpt_load(attn_ckpt, attn_ckpt_version)
+        try:
+            self.unet = UNet2DConditionModel.from_pretrained(base_ckpt, subfolder="unet").to(
+                self.device, dtype=self.weight_dtype
+            )
+            print("[INFO] UNet yüklendi")
+            
+            init_adapter(self.unet, cross_attn_cls=SkipAttnProcessor)
+            self.attn_modules = get_trainable_module(self.unet, "attention")
+            print("[INFO] Adapter başlatıldı")
+            
+            self.auto_attn_ckpt_load(attn_ckpt, attn_ckpt_version)
+            print("[INFO] Attention checkpoint yüklendi")
+        except Exception as e:
+            print(f"[ERROR] UNet yüklenirken hata: {e}")
+            raise
 
         # Compile (isteğe bağlı)
         if compile:
             try:
+                print("[INFO] Model compile ediliyor...")
                 self.unet = torch.compile(self.unet)
                 self.vae = torch.compile(self.vae, mode="reduce-overhead")
+                print("[INFO] Model compile edildi")
             except Exception as e:
                 print(f"[WARN] Compile sırasında hata: {e}. Compile edilmemiş sürüm devam ediyor.")
 
+        print("[INFO] Pipeline başarıyla başlatıldı")
+
     def auto_attn_ckpt_load(self, attn_ckpt, version):
         sub_folder = {
             "mix": "mix-48k-1024",
-            "vitonhd": "vitonhd-16k-512",
+            "vitonhd": "vitonhd-16k-512", 
             "dresscode": "dresscode-16k-512",
         }[version]
+        
         if os.path.exists(attn_ckpt):
-            load_checkpoint_in_model(self.attn_modules, os.path.join(attn_ckpt, sub_folder, "attention"))
+            checkpoint_path = os.path.join(attn_ckpt, sub_folder, "attention")
+            if os.path.exists(checkpoint_path):
+                load_checkpoint_in_model(self.attn_modules, checkpoint_path)
+                print(f"[INFO] Local checkpoint yüklendi: {checkpoint_path}")
+            else:
+                print(f"[WARN] Local checkpoint bulunamadı: {checkpoint_path}")
         else:
-            repo_path = snapshot_download(repo_id=attn_ckpt)
-            print(f"Downloaded {attn_ckpt} to {repo_path}")
-            load_checkpoint_in_model(self.attn_modules, os.path.join(repo_path, sub_folder, "attention"))
+            try:
+                repo_path = snapshot_download(repo_id=attn_ckpt)
+                checkpoint_path = os.path.join(repo_path, sub_folder, "attention")
+                load_checkpoint_in_model(self.attn_modules, checkpoint_path)
+                print(f"[INFO] Downloaded checkpoint yüklendi: {checkpoint_path}")
+            except Exception as e:
+                print(f"[ERROR] Checkpoint indirilemedi: {e}")
+                raise
 
     def run_safety_checker(self, image):
-        if self.safety_checker is None:
+        if self.safety_checker is None or self.skip_safety_check:
             has_nsfw_concept = None
         else:
-            safety_checker_input = self.feature_extractor(image, return_tensors="pt").to(self.device)
-            image, has_nsfw_concept = self.safety_checker(
-                images=image, clip_input=safety_checker_input.pixel_values.to(self.weight_dtype)
-            )
+            try:
+                safety_checker_input = self.feature_extractor(image, return_tensors="pt").to(self.device)
+                image, has_nsfw_concept = self.safety_checker(
+                    images=image, clip_input=safety_checker_input.pixel_values.to(self.weight_dtype)
+                )
+            except Exception as e:
+                print(f"[WARN] Safety checker çalıştırılamadı: {e}")
+                has_nsfw_concept = None
         return image, has_nsfw_concept
 
     def check_inputs(self, image, condition_image, mask, width, height):
@@ -220,19 +302,25 @@ class CatVTONPipeline:
         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
         image = numpy_to_pil(image)
 
-        if not self.skip_safety_check:
+        # Safety checker kontrolü
+        if not self.skip_safety_check and self.safety_checker is not None:
             current_script_directory = os.path.dirname(os.path.realpath(__file__))
             nsfw_image_path = os.path.join(os.path.dirname(current_script_directory), "resource", "img", "NSFW.jpg")
             try:
                 nsfw_image = PIL.Image.open(nsfw_image_path).resize(image[0].size)
             except Exception:
                 nsfw_image = None
-            image_np = np.array(image)
-            _, has_nsfw_concept = self.run_safety_checker(image=image_np)
-            if has_nsfw_concept is not None:
-                for i, not_safe in enumerate(has_nsfw_concept):
-                    if not_safe and nsfw_image is not None:
-                        image[i] = nsfw_image
+                
+            try:
+                image_np = np.array(image)
+                _, has_nsfw_concept = self.run_safety_checker(image=image_np)
+                if has_nsfw_concept is not None:
+                    for i, not_safe in enumerate(has_nsfw_concept):
+                        if not_safe and nsfw_image is not None:
+                            image[i] = nsfw_image
+                            print(f"[WARN] NSFW içerik tespit edildi, görüntü {i} değiştirildi")
+            except Exception as e:
+                print(f"[WARN] Safety check sırasında hata: {e}")
 
         return image