File size: 1,616 Bytes
feabc9a bafa3e3 d985a28 37c7828 d985a28 9024a26 d985a28 9024a26 d985a28 9024a26 feabc9a 9024a26 f95ef44 feabc9a 9024a26 37c7828 9024a26 f95ef44 37c7828 90933cd 9024a26 90933cd 9024a26 37c7828 5120226 9024a26 feabc9a 37c7828 9024a26 37c7828 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import torch
from diffusers import StableDiffusion3Pipeline
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
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-medium"
pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cpu")
# 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="cpu")
# Apply LoRA weights to the parameters of the model
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
# Define the path to the LoRA model (Local path in your Hugging Face Space)
lora_model_path = "./lora_model.pth" # Local path to LoRA 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]
return image
# Gradio interface
iface = gr.Interface(fn=generate_image, inputs="text", outputs="image")
iface.launch()
|