Update app.py

app.py

@@ -58,26 +58,67 @@ patch_processor_fast(SiglipProcessor)
 
 print("🌌 Loading Cosmos-Predict2 model...")
 
+# Handle authentication for gated model
+try:
+    from huggingface_hub import login
+    import os
+    
+    # Try to login with token from environment variable
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:
+        login(token=hf_token)
+        print("✅ Authenticated with Hugging Face")
+    else:
+        print("⚠️ No HF_TOKEN found, trying without authentication...")
+except Exception as e:
+    print(f"⚠️ Authentication failed: {e}")
+
 # Load the model at startup
 model_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
 
-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")
+try:
+    if COSMOS_PIPELINE_AVAILABLE:
+        print("🔄 Loading with Cosmos2TextToImagePipeline...")
+        try:
+            # Try loading with safety checker first
+            pipe = Cosmos2TextToImagePipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.bfloat16,
+                use_auth_token=True  # Use authentication token
+            )
+        except ImportError as e:
+            if "cosmos_guardrail" in str(e):
+                print("⚠️ cosmos_guardrail not available, trying without safety checker...")
+                # Try loading without safety checker
+                pipe = Cosmos2TextToImagePipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.bfloat16,
+                    use_auth_token=True,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                )
+            else:
+                raise e
+    else:
+        print("🔄 Loading with DiffusionPipeline (trust_remote_code=True)...")
+        pipe = DiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            use_auth_token=True  # Use authentication token
+        )
 
-print("✅ Cosmos-Predict2 model loaded successfully!")
+    pipe.to("cuda")
+    print("✅ Cosmos-Predict2 model loaded successfully!")
+    
+except Exception as e:
+    print(f"❌ Failed to load Cosmos model: {e}")
+    print("🔄 This is likely due to the model being gated/restricted or missing dependencies")
+    print("📝 Please check the Setup Guide for authentication instructions")
+    
+    # For demo purposes, we could fall back to a different model
+    # But for now, let's just exit gracefully
+    raise e
 
 # Default negative prompt for better quality
 DEFAULT_NEGATIVE_PROMPT = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
@@ -162,7 +203,7 @@ def create_interface():
     with gr.Blocks(title="Cosmos-Predict2 ZeroGPU", theme=gr.themes.Soft()) as interface:
         gr.Markdown("""
         # 🌌 Cosmos-Predict2 on ZeroGPU
-        **Powered by Huggingface Spaces • High-resolution generation • Fast inference**
+        **High-resolution generation • Fast inference**
         
         This Space uses ZeroGPU for efficient GPU allocation. The model is pre-loaded and ready to generate!
         """)