Update app.py

app.py

@@ -4,54 +4,42 @@ from huggingface_hub import login
 import os
 import gradio as gr
 
-# Retrieve Hugging Face token
-token = os.getenv("HF_TOKEN")
+# Retrieve the token from the environment variable
+token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
 if token:
-    login(token=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.")
+    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,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-    device_map="balanced"
-)
-
-# Enable attention slicing for reduced memory usage
-pipe.enable_attention_slicing()
+model_id = "stabilityai/stable-diffusion-3.5-medium"
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe.to("cuda")
 
 # Define the path to the LoRA model
-lora_model_path = "./lora_model.pth"
+lora_model_path = "https://huggingface.co/spaces/DonImages/Testing2/resolve/main/lora_model.pth"  # LoRA model path
 
-# Load and apply the LoRA weights
+# 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 (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)
 
-# 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]
+# 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=[
-        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 = gr.Interface(fn=generate_image, inputs="text", outputs="image")
 iface.launch()