from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL, UNet2DConditionModel
from transformers import CLIPTextModel, CLIPTokenizer
import torch
import gradio as gr
import spaces
from huggingface_hub import hf_hub_download
import os
import requests
import hashlib
from pathlib import Path
import re
import random
# Default LoRA for fallback
DEFAULT_LORA = "OedoSoldier/detail-tweaker-lora"
LORA_CACHE_DIR = "lora_cache"
def download_lora(url):
"""Download LoRA file from Civitai URL and cache it locally"""
# Create cache directory if it doesn't exist
os.makedirs(LORA_CACHE_DIR, exist_ok=True)
# Generate a filename from the URL
url_hash = hashlib.md5(url.encode()).hexdigest()
local_path = os.path.join(LORA_CACHE_DIR, f"{url_hash}.safetensors")
# If file already exists in cache, return the path
if os.path.exists(local_path):
print()
print("********** Lora Already Exists **********")
print()
return local_path
# Download the file
try:
response = requests.get(url, stream=True)
response.raise_for_status()
# Get the total file size
total_size = int(response.headers.get('content-length', 0))
# Download and save the file
with open(local_path, 'wb') as f:
if total_size == 0:
f.write(response.content)
else:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
print()
print("********** Lora Downloading Successfull **********")
print()
return local_path
except Exception as e:
print()
print(f"Error downloading LoRA: {str(e)}")
print()
return None
def is_civitai_url(url):
"""Check if the URL is a valid Civitai download URL"""
return bool(re.match(r'https?://civitai\.com/api/download/models/\d+', url))
@spaces.GPU
def generate_image(prompt, negative_prompt, lora_url, num_inference_steps=30, guidance_scale=7.0,
model="Real6.0", num_images=1, width=512, height=512,seed=None):
if model == "Real5.0":
model_id = "SG161222/Realistic_Vision_V5.0_noVAE"
elif model == "Real5.1":
model_id = "SG161222/Realistic_Vision_V5.1_noVAE"
elif model == "majicv7":
model_id = "digiplay/majicMIX_realistic_v7"
else:
model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
# Initialize models
vae = AutoencoderKL.from_pretrained(
model_id,
subfolder="vae"
).to("cuda")
text_encoder = CLIPTextModel.from_pretrained(
model_id,
subfolder="text_encoder"
).to("cuda")
tokenizer = CLIPTokenizer.from_pretrained(
model_id,
subfolder="tokenizer"
)
unet = UNet2DConditionModel.from_pretrained(
model_id,
subfolder="unet"
).to("cuda")
pipe = DiffusionPipeline.from_pretrained(
model_id,
text_encoder=text_encoder,
tokenizer=tokenizer,
vae=vae
).to("cuda")
# Load LoRA weights
try:
if lora_url and lora_url.strip():
if is_civitai_url(lora_url):
# Download and load Civitai LoRA
lora_path = download_lora(lora_url)
if lora_path:
pipe.load_lora_weights(lora_path)
print()
print("********** URL Lora Loaded **********")
print()
else:
pipe.load_lora_weights(DEFAULT_LORA)
print()
print("********** Default Lora Loaded **********")
print()
# If it's a HuggingFace repo path
elif '/' in lora_url and not lora_url.startswith('http'):
pipe.load_lora_weights(lora_url)
print()
print("********** URL Lora Loaded **********")
print()
else:
pipe.load_lora_weights(DEFAULT_LORA)
print()
print("********** Default Lora Loaded **********")
print()
else:
pipe.load_lora_weights(DEFAULT_LORA)
except Exception as e:
print()
print(f"Error loading LoRA weights: {str(e)}")
print()
pipe.load_lora_weights(DEFAULT_LORA)
if model == "Real6.0":
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
pipe.scheduler.config,
algorithm_type="dpmsolver++",
use_karras_sigmas=True
)
if seed is None:
seed = random.randint(0, 2**32 - 1)
generator = torch.manual_seed(seed)
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).to("cuda")
negative_text_inputs = tokenizer(
negative_prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).to("cuda")
prompt_embeds = text_encoder(text_inputs.input_ids)[0]
negative_prompt_embeds = text_encoder(negative_text_inputs.input_ids)[0]
# Generate the image
result = pipe(
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
cross_attention_kwargs={"scale": 1},
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
width=width,
height=height,
num_images_per_prompt=num_images,
generator=generator
)
torch.cuda.empty_cache()
return result.images,seed
def clean_lora_cache():
"""Clean the LoRA cache directory"""
if os.path.exists(LORA_CACHE_DIR):
for file in os.listdir(LORA_CACHE_DIR):
file_path = os.path.join(LORA_CACHE_DIR, file)
try:
if os.path.isfile(file_path):
os.unlink(file_path)
except Exception as e:
print(f"Error deleting {file_path}: {str(e)}")
title = """<h1 align="center">ProFaker</h1>"""
# Create the Gradio interface
with gr.Blocks() as demo:
gr.HTML(title)
with gr.Row():
with gr.Column():
# Input components
prompt = gr.Textbox(
label="Prompt",
info="Enter your image description here...",
lines=3
)
negative_prompt = gr.Textbox(
label="Negative Prompt",
info="Enter what you don't want in Image...",
lines=3
)
lora_input = gr.Textbox(
label="LoRA URL/Path",
info="Enter Civitai download URL or HuggingFace path (e.g., 'username/model-name')",
value=DEFAULT_LORA
)
clear_cache = gr.Button("Clear LoRA Cache")
generate_button = gr.Button("Generate Image")
with gr.Accordion("Advanced Options", open=False):
model = gr.Dropdown(
choices=["Real6.0","Real5.1","Real5.0","majicv7"],
value="Real6.0",
label="Model",
)
num_images = gr.Slider(
minimum=1,
maximum=4,
value=1,
step=1,
label="Number of Images to Generate"
)
width = gr.Slider(
minimum=256,
maximum=1024,
value=512,
step=64,
label="Image Width"
)
height = gr.Slider(
minimum=256,
maximum=1024,
value=512,
step=64,
label="Image Height"
)
steps_slider = gr.Slider(
minimum=1,
maximum=100,
value=30,
step=1,
label="Number of Steps"
)
guidance_slider = gr.Slider(
minimum=1,
maximum=10,
value=7.0,
step=0.5,
label="Guidance Scale"
)
seed_input = gr.Number(value=random.randint(0, 2**32 - 1), label="Seed (optional)")
with gr.Column():
# Output component
gallery = gr.Gallery(
label="Generated Images",
show_label=True,
elem_id="gallery",
columns=2,
rows=2
)
seed_display = gr.Textbox(label="Seed Used", interactive=False)
# Connect the interface to the generation function
generate_button.click(
fn=generate_image,
inputs=[prompt, negative_prompt, lora_input, steps_slider, guidance_slider,
model, num_images, width, height,seed_input],
outputs=[gallery,seed_display]
)
# Connect clear cache button
clear_cache.click(fn=clean_lora_cache)
demo.queue(max_size=10).launch(share=False)