Update app.py

app.py

@@ -2,10 +2,20 @@ import gradio as gr
 import numpy as np
 import random
 from PIL import Image
+import os
 
 import spaces
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, StableDiffusionImg2ImgPipeline
 import torch
+from huggingface_hub import login
+
+# Get token from Hugging Face Spaces secrets
+# This will use the environment variable HF_ACCESS_TOKEN which is the standard in HF Spaces
+hf_token = os.environ.get("HF_ACCESS_TOKEN")
+if hf_token:
+    login(hf_token)
+else:
+    print("Warning: HF_ACCESS_TOKEN not found in environment. Authentication may fail.")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-medium"
@@ -16,16 +26,31 @@ else:
     torch_dtype = torch.float32
 
 # For text-to-image
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler", shift=5)
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo_id, 
+    torch_dtype=torch_dtype,
+    use_auth_token=True  # This will use the token from login()
+)
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    model_repo_id, 
+    subfolder="scheduler", 
+    shift=5,
+    use_auth_token=True
+)
 pipe = pipe.to(device)
 
 # For image-to-image
 img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     model_repo_id,
-    torch_dtype=torch_dtype
+    torch_dtype=torch_dtype,
+    use_auth_token=True
+)
+img2img_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    model_repo_id, 
+    subfolder="scheduler", 
+    shift=5,
+    use_auth_token=True
 )
-img2img_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder="scheduler", shift=5)
 img2img_pipe = img2img_pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max