File size: 6,276 Bytes
7733cf5 9772569 7e583c2 9772569 ffbb513 9772569 ffbb513 ec0e4d3 bf9239a 7e583c2 9772569 bf9239a 9772569 bf9239a 9772569 bf9239a 2f2cbf0 bf9239a 67720c3 1589ab9 9772569 bf9239a 9772569 2a2d6e3 9772569 1589ab9 2a2d6e3 9772569 ffbb513 9772569 ffbb513 9772569 ffbb513 9772569 b0c6316 9a88fb8 b0c6316 |
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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
import os
import spaces
import gradio as gr
import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler
from huggingface_hub import hf_hub_download, login
from PIL import Image
import requests
from translatepy import Translator
import numpy as np
import random
import logging
# Setup logging
logging.basicConfig(level=logging.DEBUG)
# Get HF_TOKEN from environment
HF_TOKEN = os.getenv("HF_TOKEN")
login(token=HF_TOKEN)
translator = Translator()
# Constants
MODEL_NAME = "black-forest-labs/FLUX.1-dev"
LORA_WEIGHT_NAME = "MegaTronX/CivitAI_Flux_LoRA-SuicideGirlsv1"
WEIGHT_FILE = "SuicideGirlsv1.safetensors"
# CSS for hiding the footer
CSS = """
footer {
visibility: hidden;
}
"""
MAX_SEED = np.iinfo(np.int32).max
# Load the main model and LoRA weights
def load_model():
pipe = FluxPipeline.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)
pipe.load_lora_weights(LORA_WEIGHT_NAME, weight_name=WEIGHT_FILE)
pipe.fuse_lora(lora_scale=0.8)
pipe.to("cuda")
return pipe
pipe = load_model()
def clear_gallery():
return None
# Function to generate images
@spaces.GPU()
def generate_image(
prompt,
width=768,
height=1024,
scale=3.5,
steps=24,
seed=42,
nums=1,
progress=gr.Progress(track_tqdm=True)
):
# Clears the gallery
yield None, seed
if seed == 42:
seed = random.randint(0, MAX_SEED)
seed = int(seed)
generator = torch.Generator().manual_seed(seed)
prompt = str(translator.translate(prompt, 'English'))
logging.debug(f"Translated Prompt: {prompt}")
try:
images = pipe(
prompt,
width=width,
height=height,
guidance_scale=scale,
num_inference_steps=steps,
generator=generator,
output_type="pil",
max_sequence_length=512,
num_images_per_prompt=nums,
).images
logging.debug(f"Generated Images: {images}")
# Save the first image as a test
images[0].save("test_image.png")
logging.info(f"Image saved to test_image.png")
except Exception as e:
logging.error(f"Error during image generation: {e}")
return images, seed
examples = [
"SuicideGirl with black hair and green eyes",
"close up portrait, Amidst the interplay of light and shadows in a photography studio, a soft spotlight traces the contours of a face, highlighting a figure clad in a sleek black turtleneck. The garment, hugging the skin with subtle luxury, complements the Caucasian model's understated makeup, embodying minimalist elegance. Behind, a pale gray backdrop extends, its fine texture shimmering subtly in the dim light, artfully balancing the composition and focusing attention on the subject. In a palette of black, gray, and skin tones, simplicity intertwines with profundity, as every detail whispers untold stories.",
"Caucasian, The image features a young woman of European descent standing in an studio setting, surrounded by silk. (She is wearing a silk dress), paired with a bold. Her brown hair is wet and tousled, falling naturally around her face, giving her a raw and edgy look. The woman has an intense and direct gaze, adding to the dramatic feel of the image. The backdrop is flowing silk, big silk. The overall composition blends elements of fashion and nature, creating a striking and powerful visual",
"A black and white portrait of a young woman with a captivating gaze. She's bundled up in a cozy black sweater, hands gently cupped near her face. The monochromatic tones highlight her delicate features and the contemplative mood of the image",
"Fashion photography portrait, close up portrait, (a woman of European descent is surrounded by lava rock and magma from head to neck, red magma hair, wear volcanic lava rock magma outfit coat lava rock magma fashion costume with ruffled layers"
]
# Gradio Interface
with gr.Blocks(css=CSS, theme="ocean") as demo:
with gr.Group():
with gr.Row():
prompt = gr.Textbox(label='Enter Your Prompt (multilingual)', scale=6)
submit = gr.Button(scale=1, variant='primary')
img = gr.Gallery(label="Gallery", columns=1, preview=True, height=600)
with gr.Accordion("Advanced Options", open=False):
with gr.Row():
width = gr.Slider(
label="Width",
minimum=512,
maximum=1280,
step=8,
value=768,
)
height = gr.Slider(
label="Height",
minimum=512,
maximum=1280,
step=8,
value=1024,
)
with gr.Row():
scale = gr.Slider(
label="Guidance Scale",
minimum=0,
maximum=50,
step=0.1,
value=3.0,
)
steps = gr.Slider(
label="Steps",
minimum=1,
maximum=50,
step=1,
value=28,
)
with gr.Row():
seed = gr.Slider(
label="Seed (-1 Random)",
minimum=-1,
maximum=MAX_SEED,
step=1,
value=0,
visible=True
)
nums = gr.Slider(
label="Image Numbers",
minimum=1,
maximum=4,
step=1,
value=1,
scale=1,
)
gr.Examples(
examples=examples,
inputs=prompt,
outputs=[img, seed],
fn=generate_image,
cache_examples=True,
cache_mode='lazy'
)
gr.on(
triggers=[
prompt.submit,
submit.click,
],
fn=generate_image,
inputs=[
prompt,
width,
height,
scale,
steps,
seed,
nums
],
outputs=[img, seed],
api_name="run",
)
# Clears the gallery before generating a new image
submit.click(fn=clear_gallery, outputs=img, queue=False)
demo.queue().launch() |