a
This commit is contained in:
parent
3101d1ec58
commit
0a0dc79d44
426
cogs/stable_diffusion_cog.py
Normal file
426
cogs/stable_diffusion_cog.py
Normal file
@ -0,0 +1,426 @@
|
|||||||
|
import discord
|
||||||
|
from discord.ext import commands
|
||||||
|
from discord import app_commands
|
||||||
|
import torch
|
||||||
|
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
|
||||||
|
import os
|
||||||
|
import io
|
||||||
|
import time
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from typing import Optional, Literal, Dict, Any, Union
|
||||||
|
|
||||||
|
class StableDiffusionCog(commands.Cog):
|
||||||
|
def __init__(self, bot):
|
||||||
|
self.bot = bot
|
||||||
|
self.model = None
|
||||||
|
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||||
|
|
||||||
|
# Set up model directories
|
||||||
|
self.models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models")
|
||||||
|
self.illustrious_dir = os.path.join(self.models_dir, "illustrious_xl")
|
||||||
|
|
||||||
|
# Create directories if they don't exist
|
||||||
|
os.makedirs(self.models_dir, exist_ok=True)
|
||||||
|
os.makedirs(self.illustrious_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# Default to Illustrious XL if available, otherwise fallback to SD 1.5
|
||||||
|
self.model_id = self.illustrious_dir if os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")) else "runwayml/stable-diffusion-v1-5"
|
||||||
|
self.model_type = "sdxl" if self.model_id == self.illustrious_dir else "sd"
|
||||||
|
self.is_generating = False
|
||||||
|
|
||||||
|
print(f"StableDiffusionCog initialized! Using device: {self.device}")
|
||||||
|
print(f"Default model: {self.model_id} (Type: {self.model_type})")
|
||||||
|
|
||||||
|
# Check if Illustrious XL is available
|
||||||
|
if self.model_id != self.illustrious_dir:
|
||||||
|
print("Illustrious XL model not found. Using default model instead.")
|
||||||
|
print(f"To download Illustrious XL, run the download_illustrious.py script.")
|
||||||
|
|
||||||
|
async def load_model(self):
|
||||||
|
"""Load the Stable Diffusion model asynchronously"""
|
||||||
|
if self.model is not None:
|
||||||
|
return True
|
||||||
|
|
||||||
|
# This could take some time, so we run it in a thread pool
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
try:
|
||||||
|
# Check if we're loading a local model or a HuggingFace model
|
||||||
|
if os.path.isdir(self.model_id):
|
||||||
|
# Local model (like Illustrious XL)
|
||||||
|
if self.model_type == "sdxl":
|
||||||
|
print(f"Loading local SDXL model from {self.model_id}...")
|
||||||
|
self.model = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: StableDiffusionXLPipeline.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
|
||||||
|
use_safetensors=True,
|
||||||
|
variant="fp16" if self.device == "cuda" else None
|
||||||
|
).to(self.device)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(f"Loading local SD model from {self.model_id}...")
|
||||||
|
self.model = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: StableDiffusionPipeline.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
|
||||||
|
use_safetensors=True,
|
||||||
|
variant="fp16" if self.device == "cuda" else None
|
||||||
|
).to(self.device)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# HuggingFace model
|
||||||
|
if "xl" in self.model_id.lower():
|
||||||
|
self.model_type = "sdxl"
|
||||||
|
print(f"Loading SDXL model from HuggingFace: {self.model_id}...")
|
||||||
|
self.model = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: StableDiffusionXLPipeline.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
|
||||||
|
use_safetensors=True,
|
||||||
|
variant="fp16" if self.device == "cuda" else None
|
||||||
|
).to(self.device)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.model_type = "sd"
|
||||||
|
print(f"Loading SD model from HuggingFace: {self.model_id}...")
|
||||||
|
self.model = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: StableDiffusionPipeline.from_pretrained(
|
||||||
|
self.model_id,
|
||||||
|
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
|
||||||
|
).to(self.device)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Use DPM++ 2M Karras scheduler for better quality
|
||||||
|
self.model.scheduler = DPMSolverMultistepScheduler.from_config(
|
||||||
|
self.model.scheduler.config,
|
||||||
|
algorithm_type="dpmsolver++",
|
||||||
|
use_karras_sigmas=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# Enable attention slicing for lower memory usage
|
||||||
|
if hasattr(self.model, "enable_attention_slicing"):
|
||||||
|
self.model.enable_attention_slicing()
|
||||||
|
|
||||||
|
# Enable memory efficient attention if available (for SDXL)
|
||||||
|
if hasattr(self.model, "enable_xformers_memory_efficient_attention"):
|
||||||
|
try:
|
||||||
|
self.model.enable_xformers_memory_efficient_attention()
|
||||||
|
print("Enabled xformers memory efficient attention")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Could not enable xformers: {e}")
|
||||||
|
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error loading Stable Diffusion model: {e}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
return False
|
||||||
|
|
||||||
|
@app_commands.command(
|
||||||
|
name="generate",
|
||||||
|
description="Generate an image using Stable Diffusion running locally on GPU"
|
||||||
|
)
|
||||||
|
@app_commands.describe(
|
||||||
|
prompt="The text prompt to generate an image from",
|
||||||
|
negative_prompt="Things to avoid in the generated image",
|
||||||
|
steps="Number of inference steps (higher = better quality but slower)",
|
||||||
|
guidance_scale="How closely to follow the prompt (higher = more faithful but less creative)",
|
||||||
|
width="Image width (must be a multiple of 8)",
|
||||||
|
height="Image height (must be a multiple of 8)",
|
||||||
|
seed="Random seed for reproducible results (leave empty for random)",
|
||||||
|
hidden="Whether to make the response visible only to you"
|
||||||
|
)
|
||||||
|
async def generate_image(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
prompt: str,
|
||||||
|
negative_prompt: Optional[str] = None,
|
||||||
|
steps: Optional[int] = 30,
|
||||||
|
guidance_scale: Optional[float] = 7.5,
|
||||||
|
width: Optional[int] = 1024,
|
||||||
|
height: Optional[int] = 1024,
|
||||||
|
seed: Optional[int] = None,
|
||||||
|
hidden: Optional[bool] = False
|
||||||
|
):
|
||||||
|
"""Generate an image using Stable Diffusion running locally on GPU"""
|
||||||
|
# Check if already generating an image
|
||||||
|
if self.is_generating:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"⚠️ I'm already generating an image. Please wait until the current generation is complete.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Validate parameters
|
||||||
|
if steps < 1 or steps > 150:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"⚠️ Steps must be between 1 and 150.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if guidance_scale < 1 or guidance_scale > 20:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"⚠️ Guidance scale must be between 1 and 20.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if width % 8 != 0 or height % 8 != 0:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"⚠️ Width and height must be multiples of 8.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Different size limits for SDXL vs regular SD
|
||||||
|
max_size = 1536 if self.model_type == "sdxl" else 1024
|
||||||
|
min_size = 512 if self.model_type == "sdxl" else 256
|
||||||
|
|
||||||
|
if width < min_size or width > max_size or height < min_size or height > max_size:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
f"⚠️ Width and height must be between {min_size} and {max_size} for the current model type ({self.model_type.upper()}).",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Defer the response since this will take some time
|
||||||
|
await interaction.response.defer(ephemeral=hidden)
|
||||||
|
|
||||||
|
# Set the flag to indicate we're generating
|
||||||
|
self.is_generating = True
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Load the model if not already loaded
|
||||||
|
if not await self.load_model():
|
||||||
|
await interaction.followup.send(
|
||||||
|
"❌ Failed to load the Stable Diffusion model. Check the logs for details.",
|
||||||
|
ephemeral=hidden
|
||||||
|
)
|
||||||
|
self.is_generating = False
|
||||||
|
return
|
||||||
|
|
||||||
|
# Generate a random seed if none provided
|
||||||
|
if seed is None:
|
||||||
|
seed = int(time.time())
|
||||||
|
|
||||||
|
# Set the generator for reproducibility
|
||||||
|
generator = torch.Generator(device=self.device).manual_seed(seed)
|
||||||
|
|
||||||
|
# Create a status message
|
||||||
|
model_name = "Illustrious XL" if self.model_id == self.illustrious_dir else self.model_id
|
||||||
|
status_message = f"🖌️ Generating image with {model_name}\n"
|
||||||
|
status_message += f"🔤 Prompt: `{prompt}`\n"
|
||||||
|
status_message += f"📊 Parameters: Steps={steps}, CFG={guidance_scale}, Size={width}x{height}, Seed={seed}"
|
||||||
|
if negative_prompt:
|
||||||
|
status_message += f"\n🚫 Negative prompt: `{negative_prompt}`"
|
||||||
|
status_message += "\n\n⏳ Please wait, this may take a minute..."
|
||||||
|
|
||||||
|
status = await interaction.followup.send(status_message, ephemeral=hidden)
|
||||||
|
|
||||||
|
# Run the generation in a thread pool to not block the bot
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
|
||||||
|
# Different generation parameters for SDXL vs regular SD
|
||||||
|
if self.model_type == "sdxl":
|
||||||
|
# For SDXL models
|
||||||
|
image = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: self.model(
|
||||||
|
prompt=prompt,
|
||||||
|
negative_prompt=negative_prompt,
|
||||||
|
num_inference_steps=steps,
|
||||||
|
guidance_scale=guidance_scale,
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
generator=generator
|
||||||
|
).images[0]
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# For regular SD models
|
||||||
|
image = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: self.model(
|
||||||
|
prompt=prompt,
|
||||||
|
negative_prompt=negative_prompt,
|
||||||
|
num_inference_steps=steps,
|
||||||
|
guidance_scale=guidance_scale,
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
generator=generator
|
||||||
|
).images[0]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert the image to bytes for Discord upload
|
||||||
|
image_binary = io.BytesIO()
|
||||||
|
image.save(image_binary, format="PNG")
|
||||||
|
image_binary.seek(0)
|
||||||
|
|
||||||
|
# Create a file to send
|
||||||
|
file = discord.File(fp=image_binary, filename="stable_diffusion_image.png")
|
||||||
|
|
||||||
|
# Create an embed with the image and details
|
||||||
|
embed = discord.Embed(
|
||||||
|
title="🖼️ Stable Diffusion Image",
|
||||||
|
description=f"**Prompt:** {prompt}",
|
||||||
|
color=0x9C84EF
|
||||||
|
)
|
||||||
|
if negative_prompt:
|
||||||
|
embed.add_field(name="Negative Prompt", value=negative_prompt, inline=False)
|
||||||
|
|
||||||
|
# Add model info to the embed
|
||||||
|
model_info = f"Model: {model_name}\nType: {self.model_type.upper()}"
|
||||||
|
embed.add_field(name="Model", value=model_info, inline=False)
|
||||||
|
|
||||||
|
# Add generation parameters
|
||||||
|
embed.add_field(
|
||||||
|
name="Parameters",
|
||||||
|
value=f"Steps: {steps}\nGuidance Scale: {guidance_scale}\nSize: {width}x{height}\nSeed: {seed}",
|
||||||
|
inline=False
|
||||||
|
)
|
||||||
|
|
||||||
|
embed.set_image(url="attachment://stable_diffusion_image.png")
|
||||||
|
embed.set_footer(text=f"Generated by {interaction.user.display_name}", icon_url=interaction.user.display_avatar.url)
|
||||||
|
|
||||||
|
# Send the image
|
||||||
|
await interaction.followup.send(file=file, embed=embed, ephemeral=hidden)
|
||||||
|
|
||||||
|
# Try to delete the status message
|
||||||
|
try:
|
||||||
|
await status.delete()
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
await interaction.followup.send(
|
||||||
|
f"❌ Error generating image: {str(e)}",
|
||||||
|
ephemeral=hidden
|
||||||
|
)
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
finally:
|
||||||
|
# Reset the flag
|
||||||
|
self.is_generating = False
|
||||||
|
|
||||||
|
@app_commands.command(
|
||||||
|
name="sd_models",
|
||||||
|
description="List available Stable Diffusion models or change the current model"
|
||||||
|
)
|
||||||
|
@app_commands.describe(
|
||||||
|
model="The model to switch to (leave empty to just list available models)",
|
||||||
|
)
|
||||||
|
@app_commands.choices(model=[
|
||||||
|
app_commands.Choice(name="Illustrious XL (Local)", value="illustrious_xl"),
|
||||||
|
app_commands.Choice(name="Stable Diffusion 1.5", value="runwayml/stable-diffusion-v1-5"),
|
||||||
|
app_commands.Choice(name="Stable Diffusion 2.1", value="stabilityai/stable-diffusion-2-1"),
|
||||||
|
app_commands.Choice(name="Stable Diffusion XL", value="stabilityai/stable-diffusion-xl-base-1.0")
|
||||||
|
])
|
||||||
|
@commands.is_owner()
|
||||||
|
async def sd_models(
|
||||||
|
self,
|
||||||
|
interaction: discord.Interaction,
|
||||||
|
model: Optional[app_commands.Choice[str]] = None
|
||||||
|
):
|
||||||
|
"""List available Stable Diffusion models or change the current model"""
|
||||||
|
# Check if user is the bot owner
|
||||||
|
if interaction.user.id != self.bot.owner_id:
|
||||||
|
await interaction.response.send_message(
|
||||||
|
"⛔ Only the bot owner can use this command.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
if model is None:
|
||||||
|
# Just list the available models
|
||||||
|
current_model = "Illustrious XL (Local)" if self.model_id == self.illustrious_dir else self.model_id
|
||||||
|
|
||||||
|
embed = discord.Embed(
|
||||||
|
title="🤖 Available Stable Diffusion Models",
|
||||||
|
description=f"**Current model:** `{current_model}`\n**Type:** `{self.model_type.upper()}`",
|
||||||
|
color=0x9C84EF
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check if Illustrious XL is available
|
||||||
|
illustrious_status = "✅ Installed" if os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")) else "❌ Not installed"
|
||||||
|
|
||||||
|
embed.add_field(
|
||||||
|
name="Available Models",
|
||||||
|
value=(
|
||||||
|
f"• `Illustrious XL` - {illustrious_status}\n"
|
||||||
|
"• `runwayml/stable-diffusion-v1-5` - Stable Diffusion 1.5\n"
|
||||||
|
"• `stabilityai/stable-diffusion-2-1` - Stable Diffusion 2.1\n"
|
||||||
|
"• `stabilityai/stable-diffusion-xl-base-1.0` - Stable Diffusion XL"
|
||||||
|
),
|
||||||
|
inline=False
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add download instructions if Illustrious XL is not installed
|
||||||
|
if illustrious_status == "❌ Not installed":
|
||||||
|
embed.add_field(
|
||||||
|
name="Download Illustrious XL",
|
||||||
|
value=(
|
||||||
|
"To download Illustrious XL, run the `download_illustrious.py` script.\n"
|
||||||
|
"This will download the model from Civitai and set it up for use."
|
||||||
|
),
|
||||||
|
inline=False
|
||||||
|
)
|
||||||
|
|
||||||
|
embed.add_field(
|
||||||
|
name="GPU Status",
|
||||||
|
value=f"Using device: `{self.device}`\nCUDA available: `{torch.cuda.is_available()}`",
|
||||||
|
inline=False
|
||||||
|
)
|
||||||
|
if torch.cuda.is_available():
|
||||||
|
embed.add_field(
|
||||||
|
name="GPU Info",
|
||||||
|
value=f"GPU: `{torch.cuda.get_device_name(0)}`\nMemory: `{torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB`",
|
||||||
|
inline=False
|
||||||
|
)
|
||||||
|
|
||||||
|
await interaction.response.send_message(embed=embed, ephemeral=True)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Change the model
|
||||||
|
await interaction.response.defer(ephemeral=True)
|
||||||
|
|
||||||
|
# Check if we're currently generating
|
||||||
|
if self.is_generating:
|
||||||
|
await interaction.followup.send(
|
||||||
|
"⚠️ Can't change model while generating an image. Please try again later.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Unload the current model to free up VRAM
|
||||||
|
if self.model is not None:
|
||||||
|
self.model = None
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
# Set the new model ID
|
||||||
|
if model.value == "illustrious_xl":
|
||||||
|
# Check if Illustrious XL is installed
|
||||||
|
if not os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")):
|
||||||
|
await interaction.followup.send(
|
||||||
|
"❌ Illustrious XL model is not installed. Please run the `download_illustrious.py` script first.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
|
self.model_id = self.illustrious_dir
|
||||||
|
self.model_type = "sdxl"
|
||||||
|
else:
|
||||||
|
self.model_id = model.value
|
||||||
|
self.model_type = "sdxl" if "xl" in model.value.lower() else "sd"
|
||||||
|
|
||||||
|
await interaction.followup.send(
|
||||||
|
f"✅ Model changed to `{model.name}`. The model will be loaded on the next generation.",
|
||||||
|
ephemeral=True
|
||||||
|
)
|
||||||
|
|
||||||
|
async def setup(bot):
|
||||||
|
await bot.add_cog(StableDiffusionCog(bot))
|
207
download_illustrious.py
Normal file
207
download_illustrious.py
Normal file
@ -0,0 +1,207 @@
|
|||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import requests
|
||||||
|
import json
|
||||||
|
import zipfile
|
||||||
|
import shutil
|
||||||
|
import argparse
|
||||||
|
from tqdm import tqdm
|
||||||
|
import time
|
||||||
|
|
||||||
|
# Illustrious XL model information
|
||||||
|
MODEL_ID = 795765
|
||||||
|
MODEL_NAME = "Illustrious XL"
|
||||||
|
MODEL_VERSION = 1 # Version 1.0
|
||||||
|
MODEL_URL = "https://civitai.com/api/download/models/795765"
|
||||||
|
MODEL_INFO_URL = f"https://civitai.com/api/v1/models/{MODEL_ID}"
|
||||||
|
|
||||||
|
def download_file(url, destination, filename=None):
|
||||||
|
"""Download a file with progress bar"""
|
||||||
|
if filename is None:
|
||||||
|
local_filename = os.path.join(destination, url.split('/')[-1])
|
||||||
|
else:
|
||||||
|
local_filename = os.path.join(destination, filename)
|
||||||
|
|
||||||
|
with requests.get(url, stream=True) as r:
|
||||||
|
r.raise_for_status()
|
||||||
|
total_size = int(r.headers.get('content-length', 0))
|
||||||
|
|
||||||
|
# Create directory if it doesn't exist
|
||||||
|
os.makedirs(os.path.dirname(local_filename), exist_ok=True)
|
||||||
|
|
||||||
|
with open(local_filename, 'wb') as f:
|
||||||
|
with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Downloading {os.path.basename(local_filename)}") as pbar:
|
||||||
|
for chunk in r.iter_content(chunk_size=8192):
|
||||||
|
if chunk:
|
||||||
|
f.write(chunk)
|
||||||
|
pbar.update(len(chunk))
|
||||||
|
|
||||||
|
return local_filename
|
||||||
|
|
||||||
|
def create_model_index(model_dir):
|
||||||
|
"""Create a model_index.json file for the diffusers library"""
|
||||||
|
model_index = {
|
||||||
|
"_class_name": "StableDiffusionXLPipeline",
|
||||||
|
"_diffusers_version": "0.21.4",
|
||||||
|
"force_zeros_for_empty_prompt": True,
|
||||||
|
"scheduler": [
|
||||||
|
{
|
||||||
|
"_class_name": "DPMSolverMultistepScheduler",
|
||||||
|
"_diffusers_version": "0.21.4",
|
||||||
|
"beta_end": 0.012,
|
||||||
|
"beta_schedule": "scaled_linear",
|
||||||
|
"beta_start": 0.00085,
|
||||||
|
"num_train_timesteps": 1000,
|
||||||
|
"prediction_type": "epsilon",
|
||||||
|
"solver_order": 2,
|
||||||
|
"solver_type": "midpoint",
|
||||||
|
"thresholding": False,
|
||||||
|
"timestep_spacing": "leading",
|
||||||
|
"trained_betas": None,
|
||||||
|
"use_karras_sigmas": True
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"text_encoder": [
|
||||||
|
{
|
||||||
|
"_class_name": "CLIPTextModel",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"_class_name": "CLIPTextModelWithProjection",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"tokenizer": [
|
||||||
|
{
|
||||||
|
"_class_name": "CLIPTokenizer",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"_class_name": "CLIPTokenizer",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"unet": {
|
||||||
|
"_class_name": "UNet2DConditionModel",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
},
|
||||||
|
"vae": {
|
||||||
|
"_class_name": "AutoencoderKL",
|
||||||
|
"_diffusers_version": "0.21.4"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
with open(os.path.join(model_dir, "model_index.json"), "w") as f:
|
||||||
|
json.dump(model_index, f, indent=2)
|
||||||
|
|
||||||
|
def download_illustrious_xl():
|
||||||
|
"""Download and set up the Illustrious XL model"""
|
||||||
|
# Set up directories
|
||||||
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
models_dir = os.path.join(script_dir, "models")
|
||||||
|
illustrious_dir = os.path.join(models_dir, "illustrious_xl")
|
||||||
|
temp_dir = os.path.join(models_dir, "temp")
|
||||||
|
|
||||||
|
# Create directories if they don't exist
|
||||||
|
os.makedirs(models_dir, exist_ok=True)
|
||||||
|
os.makedirs(temp_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# Check if model already exists
|
||||||
|
if os.path.exists(os.path.join(illustrious_dir, "model_index.json")):
|
||||||
|
print(f"⚠️ {MODEL_NAME} model already exists at {illustrious_dir}")
|
||||||
|
choice = input("Do you want to re-download and reinstall the model? (y/n): ")
|
||||||
|
if choice.lower() != 'y':
|
||||||
|
print("Download cancelled.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Remove existing model
|
||||||
|
print(f"Removing existing {MODEL_NAME} model...")
|
||||||
|
shutil.rmtree(illustrious_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
# Create illustrious directory
|
||||||
|
os.makedirs(illustrious_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# Get model info from Civitai API
|
||||||
|
print(f"Fetching information about {MODEL_NAME} from Civitai...")
|
||||||
|
try:
|
||||||
|
response = requests.get(MODEL_INFO_URL)
|
||||||
|
response.raise_for_status()
|
||||||
|
model_info = response.json()
|
||||||
|
|
||||||
|
# Save model info for reference
|
||||||
|
with open(os.path.join(illustrious_dir, "model_info.json"), "w") as f:
|
||||||
|
json.dump(model_info, f, indent=2)
|
||||||
|
|
||||||
|
print(f"Model: {model_info['name']} by {model_info['creator']['username']}")
|
||||||
|
print(f"Description: {model_info['description'][:100]}...")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"⚠️ Failed to fetch model info: {e}")
|
||||||
|
print("Continuing with download anyway...")
|
||||||
|
|
||||||
|
# Download the model
|
||||||
|
print(f"Downloading {MODEL_NAME} from Civitai...")
|
||||||
|
try:
|
||||||
|
# Download to temp directory
|
||||||
|
model_file = download_file(MODEL_URL, temp_dir, "illustrious_xl.safetensors")
|
||||||
|
|
||||||
|
# Move the file to the model directory
|
||||||
|
print(f"Setting up {MODEL_NAME} model...")
|
||||||
|
|
||||||
|
# Create the necessary directory structure for diffusers
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "unet"), exist_ok=True)
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "vae"), exist_ok=True)
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "text_encoder"), exist_ok=True)
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "text_encoder_2"), exist_ok=True)
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "tokenizer"), exist_ok=True)
|
||||||
|
os.makedirs(os.path.join(illustrious_dir, "tokenizer_2"), exist_ok=True)
|
||||||
|
|
||||||
|
# Move the model file to the unet directory
|
||||||
|
shutil.move(model_file, os.path.join(illustrious_dir, "unet", "diffusion_pytorch_model.safetensors"))
|
||||||
|
|
||||||
|
# Create a model_index.json file
|
||||||
|
create_model_index(illustrious_dir)
|
||||||
|
|
||||||
|
# Create a README.md file with information about the model
|
||||||
|
with open(os.path.join(illustrious_dir, "README.md"), "w") as f:
|
||||||
|
f.write(f"# {MODEL_NAME}\n\n")
|
||||||
|
f.write(f"Downloaded from Civitai: https://civitai.com/models/{MODEL_ID}\n\n")
|
||||||
|
f.write("This model requires the diffusers library to use.\n")
|
||||||
|
f.write("Use the /generate command in the Discord bot to generate images with this model.\n")
|
||||||
|
|
||||||
|
print(f"✅ {MODEL_NAME} model has been downloaded and set up successfully!")
|
||||||
|
print(f"Model location: {illustrious_dir}")
|
||||||
|
print("You can now use the model with the /generate command in the Discord bot.")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"❌ Error downloading or setting up the model: {e}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
|
||||||
|
# Clean up
|
||||||
|
print("Cleaning up...")
|
||||||
|
shutil.rmtree(illustrious_dir, ignore_errors=True)
|
||||||
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
print("Download failed. Please try again later.")
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Clean up temp directory
|
||||||
|
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
parser = argparse.ArgumentParser(description=f"Download and set up the {MODEL_NAME} model from Civitai")
|
||||||
|
parser.add_argument("--force", action="store_true", help="Force download even if the model already exists")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if args.force:
|
||||||
|
# Remove existing model if it exists
|
||||||
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
illustrious_dir = os.path.join(script_dir, "models", "illustrious_xl")
|
||||||
|
if os.path.exists(illustrious_dir):
|
||||||
|
print(f"Removing existing {MODEL_NAME} model...")
|
||||||
|
shutil.rmtree(illustrious_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
download_illustrious_xl()
|
50
install_stable_diffusion.py
Normal file
50
install_stable_diffusion.py
Normal file
@ -0,0 +1,50 @@
|
|||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
|
||||||
|
def install_dependencies():
|
||||||
|
"""Install the required dependencies for Stable Diffusion."""
|
||||||
|
print("Installing Stable Diffusion dependencies...")
|
||||||
|
|
||||||
|
# List of required packages
|
||||||
|
packages = [
|
||||||
|
"torch",
|
||||||
|
"diffusers",
|
||||||
|
"transformers",
|
||||||
|
"accelerate"
|
||||||
|
]
|
||||||
|
|
||||||
|
# Check if CUDA is available
|
||||||
|
try:
|
||||||
|
import torch
|
||||||
|
cuda_available = torch.cuda.is_available()
|
||||||
|
if cuda_available:
|
||||||
|
cuda_version = torch.version.cuda
|
||||||
|
print(f"CUDA is available (version {cuda_version})")
|
||||||
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||||||
|
else:
|
||||||
|
print("CUDA is not available. Stable Diffusion will run on CPU (very slow).")
|
||||||
|
except ImportError:
|
||||||
|
print("PyTorch not installed yet. Will install with CUDA support.")
|
||||||
|
cuda_available = False
|
||||||
|
|
||||||
|
# Install each package
|
||||||
|
for package in packages:
|
||||||
|
print(f"Installing {package}...")
|
||||||
|
try:
|
||||||
|
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
|
||||||
|
print(f"Successfully installed {package}")
|
||||||
|
except subprocess.CalledProcessError as e:
|
||||||
|
print(f"Error installing {package}: {e}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
print("\nAll dependencies installed successfully!")
|
||||||
|
print("\nTo use the Stable Diffusion command:")
|
||||||
|
print("1. Restart your bot")
|
||||||
|
print("2. Use the /generate command with a text prompt")
|
||||||
|
print("3. Wait for the image to be generated (this may take some time)")
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
install_dependencies()
|
@ -28,3 +28,11 @@ google-cloud-vertexai==1.53.0
|
|||||||
protobuf==3.20.3
|
protobuf==3.20.3
|
||||||
proto-plus==1.23.0
|
proto-plus==1.23.0
|
||||||
aiosqlite
|
aiosqlite
|
||||||
|
# Stable Diffusion dependencies
|
||||||
|
torch
|
||||||
|
diffusers
|
||||||
|
transformers
|
||||||
|
accelerate
|
||||||
|
tqdm
|
||||||
|
safetensors
|
||||||
|
xformers
|
||||||
|
56
stable_diffusion_readme.md
Normal file
56
stable_diffusion_readme.md
Normal file
@ -0,0 +1,56 @@
|
|||||||
|
# Stable Diffusion Discord Bot Command
|
||||||
|
|
||||||
|
This feature adds a Stable Diffusion image generation command to your Discord bot, running locally on your GPU.
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
1. Run the installation script to install the required dependencies:
|
||||||
|
```
|
||||||
|
python install_stable_diffusion.py
|
||||||
|
```
|
||||||
|
|
||||||
|
2. Make sure you have a compatible GPU with CUDA support. The command will work on CPU but will be extremely slow.
|
||||||
|
|
||||||
|
3. Restart your bot after installing the dependencies.
|
||||||
|
|
||||||
|
## Commands
|
||||||
|
|
||||||
|
### `/generate`
|
||||||
|
Generate an image using Stable Diffusion running locally on your GPU.
|
||||||
|
|
||||||
|
**Parameters:**
|
||||||
|
- `prompt` (required): The text prompt to generate an image from
|
||||||
|
- `negative_prompt` (optional): Things to avoid in the generated image
|
||||||
|
- `steps` (optional, default: 30): Number of inference steps (higher = better quality but slower)
|
||||||
|
- `guidance_scale` (optional, default: 7.5): How closely to follow the prompt (higher = more faithful but less creative)
|
||||||
|
- `width` (optional, default: 512): Image width (must be a multiple of 8)
|
||||||
|
- `height` (optional, default: 512): Image height (must be a multiple of 8)
|
||||||
|
- `seed` (optional): Random seed for reproducible results (leave empty for random)
|
||||||
|
- `hidden` (optional, default: false): Whether to make the response visible only to you
|
||||||
|
|
||||||
|
### `/sd_models`
|
||||||
|
List available Stable Diffusion models or change the current model (owner only).
|
||||||
|
|
||||||
|
**Parameters:**
|
||||||
|
- `model` (optional): The model to switch to (leave empty to just list available models)
|
||||||
|
|
||||||
|
## Available Models
|
||||||
|
|
||||||
|
- Stable Diffusion 1.5 (`runwayml/stable-diffusion-v1-5`)
|
||||||
|
- Stable Diffusion 2.1 (`stabilityai/stable-diffusion-2-1`)
|
||||||
|
- Stable Diffusion XL (`stabilityai/stable-diffusion-xl-base-1.0`)
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
- CUDA-compatible GPU with at least 4GB VRAM (8GB+ recommended for larger images)
|
||||||
|
- Python 3.8+
|
||||||
|
- PyTorch with CUDA support
|
||||||
|
- diffusers library
|
||||||
|
- transformers library
|
||||||
|
- accelerate library
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
- If you encounter CUDA out-of-memory errors, try reducing the image dimensions or using a smaller model.
|
||||||
|
- The first generation might take longer as the model needs to be loaded into memory.
|
||||||
|
- If you're getting "CUDA not available" errors, make sure your GPU drivers are up to date and PyTorch is installed with CUDA support.
|
Loading…
x
Reference in New Issue
Block a user