Update app.py

app.py

@@ -1,72 +1,176 @@
 import gradio as gr
+import spaces
 import numpy as np
 import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
+from diffusers import StableDiffusion3Pipeline
+from huggingface_hub import login
+import os
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+# 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.")
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+# Login to Hugging Face using environment variable
+login(token=os.getenv("HF_TOKEN"))
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+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 = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
+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().manual_seed(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,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
         width=width,
         height=height,
-        generator=generator,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        generator=generator
     ).images[0]
 
     return image, seed
 
-
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    "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: 640px;
+    max-width: 580px;
 }
 """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+        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(
@@ -76,8 +180,7 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
+            run_button = gr.Button("Run", scale=0)
 
         result = gr.Image(label="Result", show_label=False)
 
@@ -86,7 +189,6 @@ with gr.Blocks(css=css) as demo:
                 label="Negative prompt",
                 max_lines=1,
                 placeholder="Enter a negative prompt",
-                visible=False,
             )
 
             seed = gr.Slider(
@@ -102,27 +204,26 @@ with gr.Blocks(css=css) as demo:
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    step=64,
+                    value=1024,
                 )
-
                 height = gr.Slider(
                     label="Height",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+                    step=64,
+                    value=1024,
                 )
 
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
                     minimum=0.0,
-                    maximum=10.0,
+                    maximum=20.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=7.5,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -130,25 +231,25 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=30,
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt])
+            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],
+        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,
-        ],
-        outputs=[result, seed],
+        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_choice],
+        outputs=[result, seed]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()