Update app.py

app.py

@@ -8,39 +8,57 @@ from diffusers import StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
 import torch
 from compel import Compel, ReturnedEmbeddingsType
+from huggingface_hub import login
+import os
+
+# Add your Hugging Face token here or set it as an environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")  # Get from environment variable
+# Or directly: HF_TOKEN = "hf_your_token_here"
+
+if HF_TOKEN:
+    login(token=HF_TOKEN)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Make sure to use torch.float16 consistently throughout the pipeline
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "votepurchase/waiREALCN_v14",
-    torch_dtype=torch.float16,
-    variant="fp16",  # Explicitly use fp16 variant
-    use_safetensors=True  # Use safetensors if available
-)
-
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-pipe.to(device)
-
-# Force all components to use the same dtype
-pipe.text_encoder.to(torch.float16)
-pipe.text_encoder_2.to(torch.float16)
-pipe.vae.to(torch.float16)
-pipe.unet.to(torch.float16)
-
-# 追加: Initialize Compel for long prompt processing
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True],
-    truncate_long_prompts=False
-)
+try:
+    # Make sure to use torch.float16 consistently throughout the pipeline
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "votepurchase/waiREALCN_v14",
+        torch_dtype=torch.float16,
+        variant="fp16",  # Explicitly use fp16 variant
+        use_safetensors=True,  # Use safetensors if available
+        use_auth_token=HF_TOKEN  # Pass token to download
+    )
+
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe.to(device)
+
+    # Force all components to use the same dtype
+    pipe.text_encoder.to(torch.float16)
+    pipe.text_encoder_2.to(torch.float16)
+    pipe.vae.to(torch.float16)
+    pipe.unet.to(torch.float16)
+
+    # Initialize Compel for long prompt processing
+    compel = Compel(
+        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        requires_pooled=[False, True],
+        truncate_long_prompts=False
+    )
+    
+    model_loaded = True
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    model_loaded = False
+    pipe = None
+    compel = None
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1216
 
-# 追加: Simple long prompt processing function
+# Simple long prompt processing function
 def process_long_prompt(prompt, negative_prompt=""):
     """Simple long prompt processing using Compel"""
     try:
@@ -52,7 +70,11 @@ def process_long_prompt(prompt, negative_prompt=""):
     
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    # 変更: Remove the 60-word limit warning and add long prompt check
+    if not model_loaded:
+        error_img = Image.new('RGB', (width, height), color=(50, 50, 50))
+        return error_img
+    
+    # Remove the 60-word limit warning and add long prompt check
     use_long_prompt = len(prompt.split()) > 60 or len(prompt) > 300
     
     if randomize_seed:
@@ -61,7 +83,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     generator = torch.Generator(device=device).manual_seed(seed)
     
     try:
-        # 追加: Try long prompt processing first if prompt is long
+        # Try long prompt processing first if prompt is long
         if use_long_prompt:
             print("Using long prompt processing...")
             conditioning, pooled = process_long_prompt(prompt, negative_prompt)
@@ -109,13 +131,15 @@ css = """
 with gr.Blocks(css=css) as demo:
 
     with gr.Column(elem_id="col-container"):
+        if not model_loaded:
+            gr.Markdown("⚠️ **Model failed to load. Please check your Hugging Face token.**")
 
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter your prompt (long prompts are automatically supported)",  # 変更: Updated placeholder
+                placeholder="Enter your prompt (long prompts are automatically supported)",
                 container=False,
             )
 
@@ -182,4 +206,4 @@ with gr.Blocks(css=css) as demo:
         outputs=[result]
     )
 
-demo.queue().launch()
+demo.queue().launch()