Update app.py

app.py

@@ -3,14 +3,15 @@ import sys
 import subprocess
 import importlib.util
 
-# Add 'safetensors' to our required packages.
+# Required packages including huggingface_hub for downloading files.
 required_packages = {
     "gradio": "gradio",
     "diffusers": "diffusers",
     "torch": "torch",
     "PIL": "pillow",
     "transformers": "transformers",
-    "safetensors": "safetensors"
+    "safetensors": "safetensors",
+    "huggingface_hub": "huggingface_hub"
 }
 
 def install_package(package_name):
@@ -29,60 +30,42 @@ import torch
 from PIL import Image
 from diffusers import StableDiffusionImg2ImgPipeline
 from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 
 def monkeypatch_lora(unet, lora_path, alpha=1.0):
     """
-    A simplistic implementation to merge LoRA weights into the UNet.
-    This function loads a LoRA weights file (in safetensors format) and applies the deltas
-    to the corresponding weights in the UNet. The logic here assumes that the LoRA keys include
-    either "lora_up" or "lora_down" and that the corresponding base weight can be obtained by replacing
-    these substrings with "weight".
-    
-    Parameters:
-      - unet: The UNet model of the diffusion pipeline.
-      - lora_path: Path (or identifier) to the LoRA weights file.
-      - alpha: A scaling factor for the LoRA weights.
+    Merge LoRA weights into the UNet model.
+    This function loads a LoRA weights file (safetensors format) and applies the deltas
+    to the corresponding base weights of the UNet.
     """
     print(f"Loading LoRA weights from: {lora_path}")
-    # Load the LoRA weights (assumed to be in safetensors format).
     lora_state = load_file(lora_path)
-    
-    # Get the current state dict of the UNet.
     unet_state = unet.state_dict()
-    
-    # Iterate over the LoRA weights and merge them.
+
     for key, delta in lora_state.items():
-        # Example mapping: if key contains "lora_up" or "lora_down", map it to a base weight key.
         if "lora_up" in key or "lora_down" in key:
-            # Derive the corresponding base key.
             base_key = key.replace("lora_up", "weight").replace("lora_down", "weight")
             if base_key in unet_state:
-                # Merge the LoRA delta scaled by alpha into the base weight.
                 unet_state[base_key] = unet_state[base_key] + delta.to(unet_state[base_key].device) * alpha
                 print(f"Applied LoRA delta for {base_key}")
             else:
                 print(f"Warning: Base weight {base_key} not found in UNet state dict.")
         else:
             print(f"Skipping key {key} as it does not appear to be a LoRA weight.")
-    
-    # Load the updated state dict into the UNet.
     unet.load_state_dict(unet_state)
     print("LoRA merging completed.")
 
 def load_model():
     """
     Load the base Stable Diffusion model and apply the FLUX.1-dev LoRA weights.
-    The base model is runwayml/stable-diffusion-v1-5 which contains all required components.
-    The FLUX.1-dev weights (LoRA) are then merged into the UNet.
+    If the LoRA weights file is not found locally, it will be downloaded from the Hugging Face Hub.
     """
-    # Base model identifier.
     base_model_id = "runwayml/stable-diffusion-v1-5"
     hf_token = os.environ.get("HF_TOKEN")
     if hf_token is None:
-        raise ValueError("HF_TOKEN environment variable is not set. "
-                         "Please set your Hugging Face token to access the gated repository.")
-    
-    # Load the base model with authentication.
+        raise ValueError("HF_TOKEN environment variable is not set. Please set it to access gated repositories.")
+
+    # Load the base model.
     pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
         base_model_id,
         torch_dtype=torch.float16,
@@ -91,16 +74,22 @@ def load_model():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     pipe = pipe.to(device)
     
-    # Path to the FLUX.1-dev LoRA weights.
-    # You can either use a local file path (e.g., "./flux_ghibsky_lora.safetensors") or
-    # download from the gated repository if permitted. Here, we assume a local file.
+    # Define the expected local path for the LoRA weights.
     lora_weights_path = "./flux_ghibsky_lora.safetensors"
+    
+    # If the file does not exist locally, attempt to download it.
     if not os.path.exists(lora_weights_path):
-        raise FileNotFoundError(f"LoRA weights file not found at {lora_weights_path}.")
+        print(f"LoRA weights file not found at {lora_weights_path}. Attempting to download from Hugging Face Hub...")
+        # Download the file from the gated repository.
+        lora_weights_path = hf_hub_download(
+            repo_id="black-forest-labs/FLUX.1-dev",
+            filename="flux_ghibsky_lora.safetensors",
+            use_auth_token=hf_token
+        )
+        print(f"Downloaded LoRA weights to {lora_weights_path}.")
     
     # Apply the LoRA weights to the UNet.
     monkeypatch_lora(pipe.unet, lora_weights_path, alpha=1.0)
-    
     print("Base model loaded and FLUX.1-dev LoRA weights merged.")
     return pipe
 
@@ -110,7 +99,7 @@ pipe = load_model()
 def transform_image(image: Image.Image, strength: float, steps: int) -> Image.Image:
     """
     Transforms the uploaded image into Ghibli-inspired art.
-    The prompt is automatically prefixed with "GHIBSKY style".
+    The prompt is prefixed with "GHIBSKY style".
     """
     prompt = (
         "GHIBSKY style, a portrait transformed into dreamy, Ghibli-inspired art, "