Update app.py

app.py

@@ -1,11 +1,10 @@
+import spaces  # Pour Zero GPU gratuit
 import torch
-import spaces
 import gradio as gr
 from diffusers import FluxPipeline
 from huggingface_hub import hf_hub_download
 import random
 
-
 # Chargement du modèle Flux.1-schnell
 model_id = "black-forest-labs/FLUX.1-schnell"
 
@@ -13,21 +12,40 @@ model_id = "black-forest-labs/FLUX.1-schnell"
 lora_repo = None
 lora_path = None
 
-def load_lora(repo_id):
+def load_lora(repo_id, subfolder=""):
     global lora_repo, lora_path
     try:
-        lora_path = hf_hub_download(repo_id=repo_id, filename="flux-lora.safetensors")
-        lora_repo = repo_id
-        return f"✅ LoRA chargé : {repo_id}"
+        # Essayer plusieurs noms de fichiers courants
+        possible_files = [
+            "lora.safetensors",
+            "flux_lora.safetensors", 
+            "flux-lora.safetensors",
+            "pytorch_lora_weights.safetensors"
+        ]
+        
+        # Si un subfolder est spécifié (ex: "anime")
+        kwargs = {"repo_id": repo_id}
+        if subfolder:
+            kwargs["subfolder"] = subfolder
+        
+        for filename in possible_files:
+            try:
+                lora_path = hf_hub_download(filename=filename, **kwargs)
+                lora_repo = repo_id
+                return f"✅ LoRA chargé : {repo_id}/{subfolder if subfolder else ''} ({filename})"
+            except:
+                continue
+        
+        return f"❌ Aucun fichier LoRA trouvé. Essayez avec un chemin spécifique."
     except Exception as e:
         return f"❌ Erreur LoRA : {str(e)}"
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=120)  # Allocation GPU pour 120s
 def generate(prompt, negative_prompt, width=1024, height=1024, steps=4, seed=-1, lora_scale=0.8):
     try:
-        # Chargement du pipe
+        # Chargement du pipe dans la fonction (obligatoire pour Zero GPU)
         pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-        pipe.enable_model_cpu_offload()
+        pipe.to("cuda")
         
         # Chargement LoRA si disponible
         if lora_repo and lora_path:
@@ -54,14 +72,19 @@ def generate(prompt, negative_prompt, width=1024, height=1024, steps=4, seed=-1,
         return None
 
 # Interface Gradio
-with gr.Blocks(title="Flux Schnell + LoRA", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Flux Schnell + LoRA") as demo:
     gr.Markdown("# 🎨 Flux.1 Schnell + LoRA\nGénérateur rapide (4 steps) avec support LoRA personnalisé")
     
     with gr.Row():
         with gr.Column(scale=1):
             lora_input = gr.Textbox(
                 label="Repo HuggingFace LoRA", 
-                placeholder="ex: XLabs-AI/flux-lora-collection",
+                placeholder="ex: stabilityai/stable-diffusion-xl-base-1.0",
+                value=""
+            )
+            subfolder_input = gr.Textbox(
+                label="Sous-dossier (optionnel)", 
+                placeholder="ex: anime (pour XLabs-AI)",
                 value=""
             )
             load_btn = gr.Button("Charger LoRA", variant="secondary")
@@ -89,7 +112,7 @@ with gr.Blocks(title="Flux Schnell + LoRA", theme=gr.themes.Soft()) as demo:
     generate_btn = gr.Button("🚀 Générer Image", variant="primary", size="lg")
     output = gr.Image(label="Résultat", type="pil")
     
-    load_btn.click(load_lora, inputs=lora_input, outputs=status)
+    load_btn.click(load_lora, inputs=[lora_input, subfolder_input], outputs=status)
     generate_btn.click(
         generate, 
         inputs=[prompt, neg_prompt, width, height, steps, seed, lora_scale_slider], 
@@ -97,4 +120,4 @@ with gr.Blocks(title="Flux Schnell + LoRA", theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(theme=gr.themes.Soft())