Update Interface

app.py

@@ -4,16 +4,33 @@ import random
 from diffusers import DiffusionPipeline
 import torch
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "Grandediw/lora_model"  # Use the fine-tuned model
+from diffusers import DiffusionPipeline
+import torch
+from huggingface_hub import login
+
+import yaml
+
+with open("config.yaml", "r") as f:
+    config = yaml.safe_load(f)
 
-# Adjust torch data type based on device
+token = config.get("huggingface_token")
+
+# Login to Hugging Face Hub
+login(token)
+
+# Model details
+device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load the model pipeline
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# Load model
+pipe = DiffusionPipeline.from_pretrained(
+    "Grandediw/lora_model",
+    torch_dtype=torch_dtype,
+    use_auth_token=True  # Enables private model access
+)
 pipe = pipe.to(device)
 
+
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024