app.py

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

@spaces.GPU
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()