Update app.py

app.py

@@ -1,37 +1,44 @@
-from fastapi import FastAPI, HTTPException
-import base64
+import torch
+from diffusers import StableDiffusion3Pipeline
+from huggingface_hub import login
 import os
-from contextlib import asynccontextmanager
-
-# Global variable to hold the LoRA weights
-lora_weights = None
-
-# Lifespan context manager to handle the application startup
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    global lora_weights
-    lora_path = "./lora_file.pth"  # Ensure the correct file name
-    if os.path.exists(lora_path):
-        with open(lora_path, "rb") as f:
-            lora_weights = base64.b64encode(f.read()).decode("utf-8")
-        print("LoRA weights loaded and preprocessed successfully.")
-    else:
-        raise HTTPException(status_code=500, detail="LoRA file not found.")
-    yield
-    # Cleanup if necessary (but not required in this case)
-
-# Initialize FastAPI app with lifespan context
-app = FastAPI(lifespan=lifespan)
-
-@app.post("/modify-prompt")
-async def modify_prompt(prompt: str):
-    global lora_weights
-    if lora_weights is None:
-        raise HTTPException(status_code=500, detail="LoRA weights not loaded.")
+import gradio as gr
+
+# Retrieve the token from the environment variable
+token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
+if token:
+    login(token=token)  # Log in with the retrieved token
+else:
+    raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")
+
+# Load the Stable Diffusion 3.5 model
+model_id = "stabilityai/stable-diffusion-3.5-large"
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # Removed torch_dtype argument
+pipe.to("cpu")  # Ensuring it runs on CPU
+
+# Define the path to the LoRA model
+lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
+
+# Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
+def load_lora_model(pipe, lora_model_path):
+    # Load the LoRA weights
+    lora_weights = torch.load(lora_model_path, map_location="cpu")
     
-    # Combine prompt with preprocessed LoRA data
-    extended_prompt = {
-        "prompt": prompt,
-        "lora": lora_weights
-    }
-    return extended_prompt
+    # Apply weights to the UNet submodule
+    for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
+        if name in lora_weights:
+            param.data += lora_weights[name]
+
+    return pipe
+
+# Load and apply the LoRA model weights
+pipe = load_lora_model(pipe, lora_model_path)
+
+# Function to generate an image from a text prompt
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Gradio interface
+iface = gr.Interface(fn=generate_image, inputs="text", outputs="image")
+iface.launch()