File size: 3,896 Bytes
c208586 ed61611 c208586 11fd78f c208586 ed61611 c208586 ed61611 c208586 ed61611 4ac394c ed61611 c208586 4ac394c c208586 7af1646 11fd78f ed61611 c208586 5f976a3 ed61611 c208586 ed61611 c208586 11fd78f c208586 f253fa6 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
import os
from safetensors import safe_open
# Detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
print(f"Using device: {device}")
# Load base model
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=dtype
)
# Try to load LoRA weights with multiple file name options
lora_loaded = False
lora_path = "./lora"
print("🔍 Looking for LoRA files...")
if os.path.exists(lora_path):
files = os.listdir(lora_path)
print(f"📁 Files in lora folder: {files}")
# List of possible LoRA filenames to try
possible_names = [
"adapter_model.safetensors",
"pytorch_lora_weights.safetensors",
"lora_weights.safetensors"
]
lora_file_found = None
for filename in possible_names:
if filename in files:
lora_file_found = filename
print(f"✅ Found LoRA file: {filename}")
break
if lora_file_found:
model_path = os.path.join(lora_path, lora_file_found)
file_size = os.path.getsize(model_path)
print(f"📊 File size: {file_size:,} bytes ({file_size/1024/1024:.2f} MB)")
# Test if the file is readable
try:
print("🔧 Testing file integrity...")
with safe_open(model_path, framework="pt") as f:
keys = list(f.keys())
print(f"✅ File is valid! Found {len(keys)} tensors")
# Now try to load it with diffusers
print("🎯 Loading with diffusers...")
if lora_file_found == "pytorch_lora_weights.safetensors":
# For pytorch_lora_weights.safetensors, try loading it directly
pipe.load_lora_weights(lora_path, weight_name=lora_file_found)
else:
# For standard naming, load normally
pipe.load_lora_weights(lora_path)
lora_loaded = True
print("✅ LoRA weights loaded successfully!")
except Exception as e:
print(f"❌ Error with {lora_file_found}: {e}")
print("💡 File appears to be corrupted or incompatible")
else:
print("❌ No LoRA files found with expected names")
else:
print("❌ LoRA folder not found")
# Move to device
pipe = pipe.to(device)
# Basic CPU optimizations
if device == "cpu":
try:
pipe.enable_attention_slicing()
print("✅ Basic CPU optimizations enabled")
except Exception as e:
print(f"⚠️ Could not enable CPU optimizations: {e}")
def generate(prompt, quality):
if quality == "Fast":
steps = 20
guidance_scale = 7.0
else: # High Quality
steps = 40
guidance_scale = 8.5
with torch.no_grad():
image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
return image
# Build Gradio UI
title = "Fine-tuned Stable Diffusion"
if lora_loaded:
title += " with LoRA ✨"
description = "✅ LoRA weights loaded! Your custom trained model is active."
else:
title += " (Base Model)"
description = "⚠️ Running with base model only. Check logs for details."
description += "\nChoose 'Fast' for quicker generation or 'High Quality' for better details."
demo = gr.Interface(
fn=generate,
inputs=[
gr.Textbox(label="Enter your prompt", placeholder="A beautiful artwork..."),
gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Mode")
],
outputs=gr.Image(label="Generated Image"),
title=title,
description=description
)
if __name__ == "__main__":
demo.launch() |