Spaces:

Gen-AI-Project
/

Art-History-Image-Generation

Sleeping

Update app.py

ed61611 verified about 1 year ago

3.9 kB

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()

	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()