Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import os | |
| import torch | |
| import random | |
| from huggingface_hub import snapshot_download | |
| from diffusers import StableDiffusionXLPipeline, AutoencoderKL | |
| from diffusers import ( | |
| EulerAncestralDiscreteScheduler, | |
| DPMSolverMultistepScheduler, | |
| DPMSolverSDEScheduler, | |
| HeunDiscreteScheduler, | |
| DDIMScheduler, | |
| LMSDiscreteScheduler, | |
| PNDMScheduler, | |
| UniPCMultistepScheduler, | |
| ) | |
| from diffusers.models.attention_processor import AttnProcessor2_0 | |
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| from transformers import AutoProcessor, AutoModelForCausalLM, pipeline | |
| import requests | |
| def download_file(url, folder_path, filename): | |
| if not os.path.exists(folder_path): | |
| os.makedirs(folder_path) | |
| file_path = os.path.join(folder_path, filename) | |
| if os.path.isfile(file_path): | |
| print(f"File already exists: {file_path}") | |
| else: | |
| response = requests.get(url, stream=True) | |
| if response.status_code == 200: | |
| with open(file_path, 'wb') as file: | |
| for chunk in response.iter_content(chunk_size=1024): | |
| file.write(chunk) | |
| print(f"File successfully downloaded and saved: {file_path}") | |
| else: | |
| print(f"Error downloading the file. Status code: {response.status_code}") | |
| ckpt_dir_pony = snapshot_download(repo_id="John6666/pony-realism-v22main-sdxl") | |
| ckpt_dir_cyber = snapshot_download(repo_id="John6666/cyberrealistic-pony-v65-sdxl") | |
| ckpt_dir_stallion = snapshot_download(repo_id="TheImposterImposters/tamePony_v25") | |
| vae_pony = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_pony, "vae"), torch_dtype=torch.float16) | |
| vae_cyber = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_cyber, "vae"), torch_dtype=torch.float16) | |
| vae_stallion = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_stallion, "vae"), torch_dtype=torch.float16) | |
| pipe_pony = StableDiffusionXLPipeline.from_pretrained( | |
| ckpt_dir_pony, vae=vae_pony, torch_dtype=torch.float16, use_safetensors=True, | |
| ) | |
| pipe_cyber = StableDiffusionXLPipeline.from_pretrained( | |
| ckpt_dir_cyber, vae=vae_cyber, torch_dtype=torch.float16, use_safetensors=True, | |
| ) | |
| pipe_stallion = StableDiffusionXLPipeline.from_pretrained( | |
| ckpt_dir_stallion, vae=vae_stallion, torch_dtype=torch.float16, use_safetensors=True, | |
| ) | |
| pipe_pony = pipe_pony.to("cuda") | |
| pipe_cyber = pipe_cyber.to("cuda") | |
| pipe_stallion = pipe_stallion.to("cuda") | |
| pipe_pony.unet.set_attn_processor(AttnProcessor2_0()) | |
| pipe_cyber.unet.set_attn_processor(AttnProcessor2_0()) | |
| pipe_stallion.unet.set_attn_processor(AttnProcessor2_0()) | |
| samplers = { | |
| "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe_pony.scheduler.config), | |
| "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe_pony.scheduler.config, use_karras_sigmas=True), | |
| "Heun": HeunDiscreteScheduler.from_config(pipe_pony.scheduler.config), | |
| "DPM++ 2M SDE Karras": DPMSolverMultistepScheduler.from_config(pipe_pony.scheduler.config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"), | |
| "DPM++ 2M": DPMSolverMultistepScheduler.from_config(pipe_pony.scheduler.config), | |
| "DDIM": DDIMScheduler.from_config(pipe_pony.scheduler.config), | |
| "LMS": LMSDiscreteScheduler.from_config(pipe_pony.scheduler.config), | |
| "PNDM": PNDMScheduler.from_config(pipe_pony.scheduler.config), | |
| "UniPC": UniPCMultistepScheduler.from_config(pipe_pony.scheduler.config), | |
| } | |
| DEFAULT_POSITIVE_PREFIX = "Score_9 score_8_up score_7_up BREAK" | |
| DEFAULT_POSITIVE_SUFFIX = "(masterpiece) very_aesthetic detailed_face cinematic footage" | |
| DEFAULT_NEGATIVE_PREFIX = "Score_1 score_2 score _3 text low_res" | |
| DEFAULT_NEGATIVE_SUFFIX = "Nsfw oversaturated crappy_art low_quality blurry bad_anatomy extra_digits fewer_digits simple_background very_displeasing watermark signature" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device) | |
| enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device) | |
| def enhance_prompt(input_prompt, model_choice): | |
| if model_choice == "Medium": | |
| result = enhancer_medium("Enhance the description: " + input_prompt) | |
| enhanced_text = result[0]['summary_text'] | |
| else: | |
| result = enhancer_long("Enhance the description: " + input_prompt) | |
| enhanced_text = result[0]['summary_text'] | |
| return enhanced_text | |
| def generate_image(model_choice, additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps, | |
| guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, use_medium_enhancer, use_long_enhancer, | |
| use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix, | |
| progress=gr.Progress(track_tqdm=True)): | |
| if model_choice == "Pony Realism v22": | |
| pipe = pipe_pony | |
| elif model_choice == "Cyber Realistic Pony v65": | |
| pipe = pipe_cyber | |
| else: | |
| pipe = pipe_stallion | |
| if use_random_seed: | |
| seed = random.randint(0, 2**32 - 1) | |
| else: | |
| seed = int(seed) | |
| pipe.scheduler = samplers[sampler] | |
| pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1) | |
| full_positive_prompt = DEFAULT_POSITIVE_PREFIX + ", " if use_positive_prefix else "" | |
| if additional_positive_prompt: | |
| enhanced_prompt = additional_positive_prompt | |
| if use_medium_enhancer: | |
| medium_enhanced = enhance_prompt(enhanced_prompt, "Medium") | |
| medium_enhanced = medium_enhanced.lower().replace('.', ',') | |
| enhanced_prompt = f"{enhanced_prompt}, {medium_enhanced}" | |
| if use_long_enhancer: | |
| long_enhanced = enhance_prompt(enhanced_prompt, "Long") | |
| long_enhanced = long_enhanced.lower().replace('.', ',') | |
| enhanced_prompt = f"{enhanced_prompt}, {long_enhanced}" | |
| full_positive_prompt += enhanced_prompt | |
| if use_positive_suffix: | |
| full_positive_prompt += f", {DEFAULT_POSITIVE_SUFFIX}" | |
| full_negative_prompt = "" | |
| if use_negative_prefix: | |
| full_negative_prompt += f"{DEFAULT_NEGATIVE_PREFIX}, " | |
| full_negative_prompt += additional_negative_prompt if additional_negative_prompt else "" | |
| if use_negative_suffix: | |
| full_negative_prompt += f", {DEFAULT_NEGATIVE_SUFFIX}" | |
| try: | |
| images = pipe( | |
| prompt=full_positive_prompt, | |
| negative_prompt=full_negative_prompt, | |
| height=height, | |
| width=width, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| num_images_per_prompt=num_images_per_prompt, | |
| generator=torch.Generator(pipe.device).manual_seed(seed) | |
| ).images | |
| return images, seed, full_positive_prompt, full_negative_prompt | |
| except Exception as e: | |
| print(f"Error during image generation: {str(e)}") | |
| import traceback | |
| traceback.print_exc() | |
| return None, seed, full_positive_prompt, full_negative_prompt | |
| with gr.Blocks(theme='bethecloud/storj_theme') as demo: | |
| gr.HTML(""" | |
| <h1 align="center">Pony Realism / Cyber Realism / Stallion Dreams</h1> | |
| <p align="center"> | |
| <a href="https://huggingface.co/John6666/pony-realism-v22main-sdxl/" target="_blank">[Pony Realism]</a> | |
| <a href="https://huggingface.co/John6666/cyberrealistic-pony-v65-sdxl" target="_blank">[Cyberrealistic Pony]</a> | |
| <a href="https://huggingface.co/TheImposterImposters/tamePony_v25" target="_blank">[Tame Pony]</a><br> | |
| <a href="https://civitai.com/models/372465/pony-realism" target="_blank">[Pony Realism civitai]</a> | |
| <a href="https://civitai.com/models/443821?modelVersionId=680915" target="_blank">[Cyberrealistic Pony civitai]</a> | |
| <a href="https://civitai.com/models/722045/tame-pony-the-authenticity-machine" target="_blank">[Tame Pony civitai]</a> | |
| <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance-Long" target="_blank">[Prompt Enhancer Long]</a> | |
| <a href="https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance" target="_blank">[Prompt Enhancer Medium]</a> | |
| </p> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| model_choice = gr.Dropdown( | |
| ["Pony Realism v22", "Cyber Realistic Pony v65", "Tame Pony v25"], | |
| label="Model Choice", | |
| value="Pony Realism v22") | |
| positive_prompt = gr.Textbox(label="Positive Prompt", placeholder="Add your positive prompt here") | |
| negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Add your negative prompt here") | |
| with gr.Accordion("Advanced settings", open=False): | |
| height = gr.Slider(512, 2048, 1024, step=64, label="Height") | |
| width = gr.Slider(512, 2048, 1024, step=64, label="Width") | |
| num_inference_steps = gr.Slider(20, 100, 30, step=1, label="Number of Inference Steps") | |
| guidance_scale = gr.Slider(1, 20, 6, step=0.1, label="Guidance Scale") | |
| num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt") | |
| use_random_seed = gr.Checkbox(label="Use Random Seed", value=True) | |
| seed = gr.Number(label="Seed", value=0, precision=0) | |
| sampler = gr.Dropdown(label="Sampler", choices=list(samplers.keys()), value="Euler a") | |
| clip_skip = gr.Slider(1, 4, 2, step=1, label="Clip skip") | |
| with gr.Accordion("Enhancers", open=False): | |
| use_medium_enhancer = gr.Checkbox(label="Use Medium Prompt Enhancer", value=False) | |
| use_long_enhancer = gr.Checkbox(label="Use Long Prompt Enhancer", value=False) | |
| generate_btn = gr.Button("Generate Image") | |
| with gr.Accordion("Prefix and Suffix Settings", open=True): | |
| use_positive_prefix = gr.Checkbox(label="Use Positive Prefix", value=True, info=f"Prefix: {DEFAULT_POSITIVE_PREFIX}") | |
| use_positive_suffix = gr.Checkbox(label="Use Positive Suffix", value=True, info=f"Suffix: {DEFAULT_POSITIVE_SUFFIX}") | |
| use_negative_prefix = gr.Checkbox(label="Use Negative Prefix", value=True, info=f"Prefix: {DEFAULT_NEGATIVE_PREFIX}") | |
| use_negative_suffix = gr.Checkbox(label="Use Negative Suffix", value=True, info=f"Suffix: {DEFAULT_NEGATIVE_SUFFIX}") | |
| with gr.Column(scale=1): | |
| output_gallery = gr.Gallery(label="Result", elem_id="gallery", show_label=False) | |
| seed_used = gr.Number(label="Seed Used") | |
| full_positive_prompt_used = gr.Textbox(label="Full Positive Prompt Used") | |
| full_negative_prompt_used = gr.Textbox(label="Full Negative Prompt Used") | |
| generate_btn.click( | |
| fn=generate_image, | |
| inputs=[ | |
| model_choice, | |
| positive_prompt, negative_prompt, height, width, num_inference_steps, | |
| guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, | |
| clip_skip, use_medium_enhancer, use_long_enhancer, | |
| use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix, | |
| ], | |
| outputs=[output_gallery, seed_used, full_positive_prompt_used, full_negative_prompt_used] | |
| ) | |
| demo.launch(debug=True, mcp_server=True) |