Update app.py

app.py

@@ -3,15 +3,17 @@ import gradio as gr
 import numpy as np
 import random
 import functools
+import os
 import torch
 from diffusers import StableDiffusion3Pipeline
 from diffusers import DiffusionPipeline
 from inference import run
 from peft import LoraConfig, get_peft_model, PeftModel
 
-model_repo_id = "stabilityai/stable-diffusion-3.5-large"
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.bfloat16)
-pipe = pipe.to('cuda')
+pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large",
+                                                torch_dtype=torch.bfloat16,
+                                                token=os.environ["TOKEN"])
+pipe = pipe.to("cuda")
 
 distill_check = 'yresearch/swd-large-6-steps'
 pipe.transformer = PeftModel.from_pretrained(