Update app.py

app.py

@@ -16,12 +16,21 @@ subprocess.run(
 import gradio as gr
 import spaces
 import torch
-from diffusers import Cosmos2TextToImagePipeline
 from transformers import AutoModelForCausalLM, SiglipProcessor
 import random
 import gc
 import warnings
 
+# Try to import Cosmos-specific pipeline, fall back to generic if not available
+try:
+    from diffusers import Cosmos2TextToImagePipeline
+    COSMOS_PIPELINE_AVAILABLE = True
+    print("✅ Cosmos2TextToImagePipeline available")
+except ImportError:
+    from diffusers import DiffusionPipeline
+    COSMOS_PIPELINE_AVAILABLE = False
+    print("⚠️ Cosmos2TextToImagePipeline not available, using DiffusionPipeline with trust_remote_code")
+
 # Suppress warnings for cleaner output
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
@@ -51,10 +60,21 @@ print("🌌 Loading Cosmos-Predict2 model...")
 
 # Load the model at startup
 model_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
-pipe = Cosmos2TextToImagePipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16
-)
+
+if COSMOS_PIPELINE_AVAILABLE:
+    print("🔄 Loading with Cosmos2TextToImagePipeline...")
+    pipe = Cosmos2TextToImagePipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16
+    )
+else:
+    print("🔄 Loading with DiffusionPipeline (trust_remote_code=True)...")
+    pipe = DiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True
+    )
+
 pipe.to("cuda")
 
 print("✅ Cosmos-Predict2 model loaded successfully!")