import os import random import uuid from typing import Tuple import gradio as gr import numpy as np from PIL import Image import spaces import torch from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler from huggingface_hub import login # Log in to Hugging Face using the provided token hf_token = os.getenv("HF_TOKEN") login(hf_token) DESCRIPTIONz = """## STABLE IMAGINE 🍺""" def save_image(img): unique_name = str(uuid.uuid4()) + ".png" img.save(unique_name) return unique_name def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: if randomize_seed: seed = random.randint(0, MAX_SEED) return seed MAX_SEED = np.iinfo(np.int32).max DESCRIPTIONz = "" if not torch.cuda.is_available(): DESCRIPTIONz += """

⚠️Running on CPU, This may not work on CPU. If it runs for an extended time or if you encounter errors, try running it on a GPU by duplicating the space using @spaces.GPU(). πŸ“

""" USE_TORCH_COMPILE = 0 ENABLE_CPU_OFFLOAD = 0 if torch.cuda.is_available(): pipe = StableDiffusionXLPipeline.from_pretrained( "SG161222/RealVisXL_V4.0_Lightning", torch_dtype=torch.float16, use_safetensors=True, ) pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) pipe.to("cuda") if USE_TORCH_COMPILE: pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) else: # If CUDA is not available, fall back to CPU (not ideal for SDXL) pipe = StableDiffusionXLPipeline.from_pretrained( "SG161222/RealVisXL_V4.0_Lightning", torch_dtype=torch.float32, # safer for CPU use_safetensors=True, ) pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) pipe.to("cpu") if ENABLE_CPU_OFFLOAD: # Optionally offload to CPU with accelerate or similar, if set up pipe.enable_model_cpu_offload() LORA_OPTIONS = { "Realism (face/character)πŸ‘¦πŸ»": ("prithivMLmods/Canopus-Realism-LoRA", "Canopus-Realism-LoRA.safetensors", "rlms"), "Pixar (art/toons)πŸ™€": ("prithivMLmods/Canopus-Pixar-Art", "Canopus-Pixar-Art.safetensors", "pixar"), "Interior Architecture (house/hotel)🏠": ("prithivMLmods/Canopus-Interior-Architecture-0.1", "Canopus-Interior-Architecture-0.1Ξ΄.safetensors", "arch"), "Fashion Product (wearing/usable)πŸ‘œ": ("prithivMLmods/Canopus-Fashion-Product-Dilation", "Canopus-Fashion-Product-Dilation.safetensors", "fashion"), "Minimalistic Image (minimal/detailed)🏞️": ("prithivMLmods/Pegasi-Minimalist-Image-Style", "Pegasi-Minimalist-Image-Style.safetensors", "minimalist"), "Modern Clothing (trend/new)πŸ‘•": ("prithivMLmods/Canopus-Modern-Clothing-Design", "Canopus-Modern-Clothing-Design.safetensors", "mdrnclth"), "Animaliea (farm/wild)🫎": ("prithivMLmods/Canopus-Animaliea-Artism", "Canopus-Animaliea-Artism.safetensors", "Animaliea"), "Canes Cars (realistic/futurecars)🚘": ("prithivMLmods/Canes-Cars-Model-LoRA", "Canes-Cars-Model-LoRA.safetensors", "car"), "Art Minimalistic (paint/semireal)🎨": ("prithivMLmods/Canopus-Art-Medium-LoRA", "Canopus-Art-Medium-LoRA.safetensors", "mdm"), } for model_name, weight_name, adapter_name in LORA_OPTIONS.values(): pipe.load_lora_weights(model_name, weight_name=weight_name, adapter_name=adapter_name) pipe.to("cuda") style_list = [ { "name": "3840 x 2160", "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic", "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly", }, # Add more style dicts here if needed ] styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list} def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]: if style_name in styles: p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME]) else: p, n = styles[DEFAULT_STYLE_NAME] if not negative: negative = "" return p.replace("{prompt}", positive), n + negative DEFAULT_STYLE_NAME = "3840 x 2160" @spaces.GPU(duration=60, enable_queue=True) def generate( prompt: str, negative_prompt: str = "", use_negative_prompt: bool = False, seed: int = 0, width: int = 1024, height: int = 1024, guidance_scale: float = 3, randomize_seed: bool = False, style_name: str = DEFAULT_STYLE_NAME, lora_model: str = "Realism (face/character)πŸ‘¦πŸ»", progress=gr.Progress(track_tqdm=True), ): seed = int(randomize_seed_fn(seed, randomize_seed)) positive_prompt, effective_negative_prompt = apply_style(style_name, prompt, negative_prompt) if not use_negative_prompt: effective_negative_prompt = "" # type: ignore model_name, weight_name, adapter_name = LORA_OPTIONS[lora_model] pipe.set_adapters(adapter_name) images = pipe( prompt=positive_prompt, negative_prompt=effective_negative_prompt, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=20, num_images_per_prompt=1, cross_attention_kwargs={"scale": 0.65}, output_type="pil", ).images image_paths = [save_image(img) for img in images] return image_paths, seed with gr.Blocks() as demo: gr.Markdown(DESCRIPTIONz) with gr.Row(): input_prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt", lines=2) use_negative_prompt = gr.Checkbox(label="Use negative prompt?", value=False) negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt", lines=2) with gr.Row(): randomize_seed = gr.Checkbox(label="Randomize Seed", value=False) seed = gr.Number(value=0, label="Seed") with gr.Row(): style_dropdown = gr.Dropdown(label="Image Style", choices=list(styles.keys()), value=DEFAULT_STYLE_NAME) lora_dropdown = gr.Dropdown(label="LoRA Model", choices=list(LORA_OPTIONS.keys()), value="Realism (face/character)πŸ‘¦πŸ»") with gr.Row(): width = gr.Slider(512, 2048, value=1024, step=64, label="Width") height = gr.Slider(512, 2048, value=1024, step=64, label="Height") with gr.Row(): guidance_scale = gr.Slider(1.0, 15.0, value=3, step=0.5, label="Guidance Scale") output_gallery = gr.Gallery(label="Generated Images", columns=[2], height="auto") output_seed = gr.Number(label="Final Seed", interactive=False) generate_button = gr.Button("Generate Images") generate_button.click( fn=generate, inputs=[ input_prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, randomize_seed, style_dropdown, lora_dropdown, ], outputs=[output_gallery, output_seed], ) demo.launch()