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 = "stabilityai/stable-diffusion-3.5-large"
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 = [
    "Samurai girl in the snow forest. Show on Sony camera f1.2",
    "A stylish Japanese woman in her early 20s stands confidently in front of a cold, industrial background in a cinematic close-up. She wears round black sunglasses, a wide-brimmed black hat, and a brown suede coat with white shearling lining layered over a black turtleneck. A bold silver chain necklace adds a sharp urban edge to her look. Her expression is fierce and composed, staring straight into the camera with quiet intensity. The lighting is cool and bluish, creating a moody, neo-noir vibe — evoking the feel of a modern Tokyo underground fashion scene.",
    "A young Indonesian woman from Bandung walks directly toward the camera across a cracked, sun-bleached desert highway under a vast, empty sky. She wears a bold futuristic high-fashion outfit: a structured matte-black coat with exaggerated shoulders and a high collar, flowing just above her minimalistic sand-toned boots. Her hijab is styled tightly and sleek under the high collar, blending seamlessly with the look — matte black with wind-swept form, no hair visible. Her makeup is bold and geometric, featuring sharp silvery eyeliner and face highlights under the eyes, catching the sun's glare. Her expression is cold, focused, and defiant — one brow slightly raised as she pierces the lens with calm intensity. The directional sunlight from the right casts elongated, sharp-edged shadows across the cracked highway. The background is pure dystopian silence: endless dry plains stretch into the distance with a few collapsed, rusted billboards leaning in the far horizon. No buildings, no people — just wind, fashion, and desert solitude. Shot in ultra high-resolution, harsh sunlight, cinematic composition, dystopian fashion editorial."
]

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"""
        # Stable Diffusion 3.5 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:**
        - **SD3.5 Photorealistic**: Specialized for photorealistic portraits and scenes
        - **Face Helper SDXL**: Enhances facial features and expressions  
        - **LCM LoRA SDXL**: Reduces inference steps for faster generation
        
        Powered by [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()