Update app.py

app.py

@@ -4,42 +4,55 @@ from huggingface_hub import login
 import os
 import gradio as gr
 
-# Retrieve the token from the environment variable
-token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
+# Retrieve Hugging Face token
+token = os.getenv("HF_TOKEN")
 if token:
-    login(token=token)  # Log in with the retrieved token
+    login(token=token)
 else:
-    raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")
+    raise ValueError("Hugging Face token not found. Please set it as a repository secret.")
 
 # Load the Stable Diffusion 3.5 model
 model_id = "stabilityai/stable-diffusion-3.5-large"
-pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe.to("cuda")
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    revision="fp16",
+    low_cpu_mem_usage=True,
+    device_map="auto"  # Automatically allocate model components to available devices
+)
+
+# Enable attention slicing for reduced memory usage
+pipe.enable_attention_slicing()
 
 # Define the path to the LoRA model
-lora_model_path = "https://huggingface.co/spaces/DonImages/Testing2/resolve/main/lora_model.pth"  # LoRA model path
+lora_model_path = "./lora_model.pth"
 
-# Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
+# Load and apply the LoRA weights
 def load_lora_model(pipe, lora_model_path):
-    # Load the LoRA weights (assuming it's a PyTorch .pth file)
     lora_weights = torch.load(lora_model_path, map_location="cuda")
+    pipe.unet.load_attn_procs(lora_weights)
+    return pipe
 
-    # Modify this section based on how LoRA is intended to interact with your Stable Diffusion model
-    # Here, we just load the weights into the model's parameters (this is a conceptual approach)
-    for name, param in pipe.named_parameters():
-        if name in lora_weights:
-            param.data += lora_weights[name]  # Apply LoRA weights to the parameters
-
-    return pipe  # Return the updated model
-
-# 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]
+# Generate image function
+def generate_image(prompt, steps, scale):
+    with torch.inference_mode():  # Avoid gradient computation for inference
+        image = pipe(prompt, num_inference_steps=steps, guidance_scale=scale).images[0]
     return image
 
 # Gradio interface
-iface = gr.Interface(fn=generate_image, inputs="text", outputs="image")
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Enter your prompt"),
+        gr.Slider(10, 50, step=1, value=30, label="Number of Inference Steps"),
+        gr.Slider(1.0, 20.0, step=0.5, value=7.5, label="Guidance Scale"),
+    ],
+    outputs="image",
+    title="Optimized Stable Diffusion with LoRA",
+    description="Generate images using Stable Diffusion 3.5 with optimized memory usage."
+)
+
+# Launch the Gradio interface
 iface.launch()