Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import spaces | |
| import numpy as np | |
| import random | |
| import torch | |
| from diffusers import StableDiffusion3Pipeline | |
| from huggingface_hub import login | |
| import os | |
| # Add this import to fix BaseTunerLayer error | |
| try: | |
| from peft.tuners.tuners_utils import BaseTunerLayer | |
| except ImportError: | |
| print("Warning: peft not installed. LoRA functionality may be limited.") | |
| # Login to Hugging Face using environment variable | |
| login(token=os.getenv("HF_TOKEN")) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.float16 if device == "cuda" else torch.float32 | |
| # Base model | |
| repo = "mojitocup/realistic-xl" | |
| pipe = StableDiffusion3Pipeline.from_pretrained( | |
| repo, | |
| torch_dtype=dtype, | |
| use_safetensors=True, | |
| variant="fp16" if dtype == torch.float16 else None | |
| ).to(device) | |
| # List of LoRA models (can expand later) | |
| loras = { | |
| "None": None, | |
| "SD3.5 Photorealistic": "prithivMLmods/SD3.5-Large-Photorealistic-LoRA", | |
| "Face Helper SDXL": "ostris/face-helper-sdxl-lora", | |
| "LCM LoRA SDXL": "latent-consistency/lcm-lora-sdxl" | |
| } | |
| MAX_SEED = np.iinfo(np.int32).max | |
| MAX_IMAGE_SIZE = 1536 | |
| class LoRAManager: | |
| """Manages LoRA loading and unloading with proper error handling""" | |
| def __init__(self, pipe): | |
| self.pipe = pipe | |
| self.current_lora = None | |
| def load_lora(self, lora_repo, lora_scale=0.8): | |
| """Load a LoRA adapter with error handling""" | |
| try: | |
| # First try to unfuse any existing LoRA | |
| self.unfuse_current_lora() | |
| # Try different common LoRA weight file names | |
| weight_names_to_try = [ | |
| "pytorch_lora_weights.safetensors", | |
| "Photorealistic-SD3.5-Large-LoRA.safetensors", # For prithivMLmods model | |
| "diffusion_pytorch_model.safetensors", | |
| None # Let diffusers auto-detect | |
| ] | |
| success = False | |
| for weight_name in weight_names_to_try: | |
| try: | |
| if weight_name: | |
| self.pipe.load_lora_weights(lora_repo, weight_name=weight_name) | |
| else: | |
| self.pipe.load_lora_weights(lora_repo) | |
| success = True | |
| break | |
| except Exception as e: | |
| print(f"Failed to load with weight_name='{weight_name}': {e}") | |
| continue | |
| if not success: | |
| print(f"Error loading LoRA {lora_repo}: No compatible weight file found") | |
| return False | |
| self.pipe.fuse_lora(lora_scale=lora_scale) | |
| self.current_lora = lora_repo | |
| print(f"Successfully loaded LoRA: {lora_repo}") | |
| return True | |
| except Exception as e: | |
| print(f"Error loading LoRA {lora_repo}: {e}") | |
| return False | |
| def unfuse_current_lora(self): | |
| """Safely unfuse current LoRA""" | |
| if self.current_lora is None: | |
| return | |
| try: | |
| self.pipe.unfuse_lora() | |
| print(f"Unfused LoRA: {self.current_lora}") | |
| self.current_lora = None | |
| except Exception as e: | |
| print(f"Warning: Could not unfuse LoRA: {e}") | |
| self.current_lora = None # Reset anyway | |
| def truncate_prompt(prompt, max_length=77): | |
| """Truncate prompt to fit CLIP token limit""" | |
| if not prompt: | |
| return prompt | |
| # Simple word-based truncation (not perfect but helps) | |
| words = prompt.split() | |
| if len(words) <= max_length: | |
| return prompt | |
| truncated = " ".join(words[:max_length]) | |
| print(f"Warning: Prompt truncated from {len(words)} to {max_length} words") | |
| return truncated | |
| def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_choice, progress=gr.Progress(track_tqdm=True)): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator(device=device).manual_seed(seed) | |
| # Truncate prompts to avoid CLIP token limit | |
| prompt = truncate_prompt(prompt, max_length=70) | |
| negative_prompt = truncate_prompt(negative_prompt, max_length=70) | |
| # Handle LoRA loading with better error handling | |
| if lora_choice != "None": | |
| lora_manager = LoRAManager(pipe) | |
| if not lora_manager.load_lora(loras[lora_choice]): | |
| raise gr.Error(f"Failed to load LoRA adapter: {lora_choice}") | |
| else: | |
| lora_manager = LoRAManager(pipe) | |
| lora_manager.unfuse_current_lora() | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| width=width, | |
| height=height, | |
| guidance_scale=guidance_scale, | |
| num_inference_steps=num_inference_steps, | |
| generator=generator | |
| ).images[0] | |
| return image, seed | |
| examples = [ | |
| "Cinematic photograpy of a young woman delicately holding a vintage vinyl record with both hands. She gazes intently through the record's central hole, as if peering into a mysterious, alternate world where you can se only her eye. The scene is illuminated by soft, warm ambient light, enhancing the nostalgic and introspective mood. Emphasize the intricate textures of the vinyl and the subtle details in her expression, with a shallow depth of field to softly blur the background and focus on her captivating presence", | |
| "A stunning blonde woman with loose, beachy waves lying on her stomach on a sandy beach. She has beautiful eyes, light freckles across her nose and cheeks, and a confident expression. Shes playfully holding a lollipop near her lips with one finger raised in a shhh gesture. She has a large black rose tattoo sleeve on one arm and is wearing a light-colored bikini top. The scene is sunlit and warm, with a softly blurred ocean in the background. The overall mood is sultry, carefree, and summery, capturing an Instagram-style beach aesthetic.", | |
| "extreme close up portrait of a woman, she has blue hair, she has a lot of tattoos,with a neutral expression that conveys a sense of calm and minimalism.,a real detailed skin, neutral white background, dynamic pose" | |
| ] | |
| css = """ | |
| #col-container { | |
| margin: 0 auto; | |
| max-width: 580px; | |
| } | |
| """ | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown(f""" | |
| # Customized Stable Diffusion Realistic XL with LoRA + Photoreal Enhancements | |
| Choose a high-quality LoRA model to enhance your generations. All models are tested and compatible with SD3.5. | |
| **Available LoRA Models:** | |
| - **Photorealistic**: Specialized for photorealistic portraits and scenes | |
| - **Face Helper**: Enhances facial features and expressions | |
| - **LCM LoRA**: Reduces inference steps for faster generation | |
| Based on [StabilityAI SD3.5 Large](https://huggingface.co/stabilityai/stable-diffusion-3.5-large). | |
| """) | |
| with gr.Row(): | |
| prompt = gr.Text( | |
| label="Prompt", | |
| show_label=False, | |
| max_lines=1, | |
| placeholder="Enter your prompt", | |
| container=False, | |
| ) | |
| run_button = gr.Button("Run", scale=0) | |
| result = gr.Image(label="Result", show_label=False) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| negative_prompt = gr.Text( | |
| label="Negative prompt", | |
| max_lines=1, | |
| placeholder="Enter a negative prompt", | |
| ) | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=0, | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| with gr.Row(): | |
| width = gr.Slider( | |
| label="Width", | |
| minimum=512, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=64, | |
| value=1024, | |
| ) | |
| height = gr.Slider( | |
| label="Height", | |
| minimum=512, | |
| maximum=MAX_IMAGE_SIZE, | |
| step=64, | |
| value=1024, | |
| ) | |
| with gr.Row(): | |
| guidance_scale = gr.Slider( | |
| label="Guidance scale", | |
| minimum=0.0, | |
| maximum=20.0, | |
| step=0.1, | |
| value=7.5, | |
| ) | |
| num_inference_steps = gr.Slider( | |
| label="Number of inference steps", | |
| minimum=1, | |
| maximum=50, | |
| step=1, | |
| value=30, | |
| ) | |
| lora_choice = gr.Dropdown( | |
| label="LoRA adapter", | |
| choices=list(loras.keys()), | |
| value="None" | |
| ) | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[prompt] | |
| ) | |
| gr.on( | |
| triggers=[run_button.click, prompt.submit, negative_prompt.submit], | |
| fn=infer, | |
| inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_choice], | |
| outputs=[result, seed] | |
| ) | |
| demo.launch() | |