Update app.py

app.py

@@ -1,53 +1,55 @@
-import os
 from huggingface_hub import InferenceClient
+import os
 import base64
 import json
 
 class LoRAInferenceWrapper:
-    def __init__(self, model_id, lora_file_path):
-        """
-        Initialize the LoRA Inference Wrapper.
-        
-        Args:
-            model_id (str): Hugging Face model ID (e.g., stabilityai/stable-diffusion-3.5-large).
-            lora_file_path (str): Path to the local LoRA file.
-        """
-        # Retrieve the API token from environment variables
-        token = os.getenv("HF_TOKEN")
-        if not token:
-            raise ValueError("HF_TOKEN is not set. Add it as a secret in the Hugging Face Space.")
+    def __init__(self, model_id, token):
+        # Initialize the InferenceClient
         self.client = InferenceClient(model_id, token=token)
-        self.lora_file_path = lora_file_path
 
     def load_lora_weights(self):
-        # Load LoRA weights (same as before)
-        try:
-            with open(self.lora_file_path, "rb") as f:
-                lora_weights = f.read()
-            return base64.b64encode(lora_weights).decode('utf-8')
-        except FileNotFoundError:
-            raise Exception(f"LoRA file not found at path: {self.lora_file_path}")
+        # Define the path to the LoRA model
+        lora_model_path = "./lora_file.pth"  # Assuming the file is saved locally
+
+        # Check if the file exists
+        if not os.path.exists(lora_model_path):
+            raise FileNotFoundError(f"LoRA file not found at path: {lora_model_path}")
+
+        # Load the LoRA weights from the local file
+        with open(lora_model_path, "rb") as f:
+            return f.read()  # Return the raw bytes of the LoRA file
+
+    def preprocess_lora_weights(self, lora_weights):
+        # Preprocess the LoRA weights (e.g., Base64 encoding for JSON compatibility)
+        return base64.b64encode(lora_weights).decode("utf-8")
 
     def generate_with_lora(self, prompt):
-        # Generate image (same as before)
-        processed_lora = self.load_lora_weights()
+        # Load and preprocess the LoRA weights
+        lora_weights = self.load_lora_weights()
+        processed_lora = self.preprocess_lora_weights(lora_weights)
+
+        # Combine the prompt and LoRA data as a single input
         extended_prompt = json.dumps({
             "prompt": prompt,
             "lora": processed_lora
         })
+
+        # Generate the output using the InferenceClient
         result = self.client.text_to_image(prompt=extended_prompt)
         return result
 
 # Example usage
-if __name__ == "__main__":
-    model_id = "stabilityai/stable-diffusion-3.5-large"
-    lora_file_path = "./lora_file.pth"
-
-    lora_client = LoRAInferenceWrapper(model_id, lora_file_path)
-
-    prompt = "A futuristic city skyline with neon lights."
-    try:
-        result = lora_client.generate_with_lora(prompt)
-        print("Generated image result:", result)
-    except Exception as e:
-        print("Error:", str(e))
+model_id = "stabilityai/stable-diffusion-3.5-large"
+token = "hf_YOUR_HF_API_TOKEN"  # Replace with your Hugging Face token
+
+# Initialize the wrapper
+lora_client = LoRAInferenceWrapper(model_id, token)
+
+# Generate an image with the LoRA file applied
+prompt = "The same woman, smiling at the beach."
+try:
+    result = lora_client.generate_with_lora(prompt)
+    print("Generated image:", result)
+except Exception as e:
+    print("Error:", str(e))