Update app.py

app.py

@@ -120,7 +120,8 @@ print("Using device:", device)
 # ============================================================
 
 #AIO_REPO_ID = "Pr0f3ssi0n4ln00b/Phr00t-Qwen-Rapid-AIO"
-AIO_REPO_ID = "Phr00t/Qwen-Image-Edit-Rapid-AIO"
+#AIO_REPO_ID = "Phr00t/Qwen-Image-Edit-Rapid-AIO"
+AIO_REPO_ID = "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19"
 DEFAULT_AIO_VERSION = "v19"
 
 _VER_RE = re.compile(r"^v\d+$")
@@ -164,43 +165,26 @@ dtype = torch.bfloat16
 
 
 def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline:
-    # Clean version name
-    ver = str(version).strip().lower()
-    if ver.startswith("v"):
-        ver = ver[1:]
-    
-    subfolder = f"v{ver}"   # e.g. "v19"
-    
-    print(f"📦 Loading AIO transformer: {AIO_REPO_ID} / {subfolder}")
+    print(f"📦 Loading AIO transformer from: {AIO_REPO_ID}")
 
     try:
-        # Load the transformer using the proper subfolder from the original repo
         transformer = QwenImageTransformer2DModel.from_pretrained(
             AIO_REPO_ID,
-            subfolder=subfolder,
             torch_dtype=dtype,
             device_map="cuda",
         )
-        print(f"✅ Successfully loaded transformer from {subfolder}")
+        print("✅ Successfully loaded the Rapid AIO transformer!")
     except Exception as e:
-        print(f"❌ Failed to load from subfolder {subfolder}: {e}")
-        print("Trying root of repo as fallback...")
-        # Fallback: try loading without subfolder (sometimes weights are at root)
-        transformer = QwenImageTransformer2DModel.from_pretrained(
-            AIO_REPO_ID,
-            torch_dtype=dtype,
-            device_map="cuda",
-        )
-    
-    # Now build the full pipeline
+        print(f"❌ Failed to load transformer: {e}")
+        raise  # Stop here so we see the clear error
+
     p = QwenImageEditPlusPipeline.from_pretrained(
         "Qwen/Qwen-Image-Edit-2511",
         transformer=transformer,
         torch_dtype=dtype,
     ).to(device)
-    
-    return p
 
+    return p
 
 # Forgiving load: try env/default version, fallback to v19 if it fails
 try: