Update app.py

app.py

@@ -1,127 +1,62 @@
-import spaces  # Pour Zero GPU gratuit
+import spaces
 import torch
 import gradio as gr
 from diffusers import FluxPipeline
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, login
 import random
 import os
 
-# Authentification HuggingFace
-from huggingface_hub import login
+# Authentification (gardez ça)
 hf_token = os.getenv("HF_TOKEN")
 print(f"Token trouvé : {bool(hf_token)}")
-print(f"Longueur du token : {len(hf_token) if hf_token else 0}")
-
-
 if hf_token:
     login(token=hf_token)
-    print("✅ Authentifié sur HuggingFace")
-else:
-    print("⚠️ HF_TOKEN non trouvé - certains modèles peuvent être inaccessibles")
-
-
+    print("✅ Authentifié")
 
-
-# Chargement du modèle Flux.1-schnell
 model_id = "black-forest-labs/FLUX.1-schnell"
-
-# Variables globales pour LoRA
 lora_repo = None
 lora_path = None
 
 def load_lora(repo_id, style):
     global lora_repo, lora_path
     try:
-        # Construire le nom du fichier basé sur le style
         filename = f"{style}_lora.safetensors"
-        
         lora_path = hf_hub_download(repo_id=repo_id, filename=filename)
         lora_repo = repo_id
-        return f"✅ LoRA chargé : {repo_id} ({filename})"
+        return f"✅ LoRA: {repo_id} ({filename})"
     except Exception as e:
-        return f"❌ Erreur LoRA : {str(e)}"
+        return f"❌ {e}"
 
-@spaces.GPU(duration=120)  # Allocation GPU pour 120s
+@spaces.GPU(duration=120)
 def generate(prompt, negative_prompt, width=1024, height=1024, steps=4, seed=-1, lora_scale=0.8):
-    try:
-        # Chargement du pipe dans la fonction (obligatoire pour Zero GPU)
-        pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-        pipe.to("cuda")
-        
-        # Chargement LoRA si disponible
-        if lora_repo and lora_path:
-            pipe.load_lora_weights(lora_path)
-            pipe.fuse_lora(lora_scale=lora_scale)
-        
-        generator = torch.Generator("cuda").manual_seed(seed if seed != -1 else random.randint(0, 2**32))
-        
-        image = pipe(
-            prompt, 
-            negative_prompt=negative_prompt,
-            height=height, 
-            width=width, 
-            num_inference_steps=steps,
-            guidance_scale=0.0, 
-            generator=generator
-        ).images[0]
-        
-        # Nettoyage VRAM
-        del pipe
-        torch.cuda.empty_cache()
-        return image
-    except Exception as e:
-        print(f"Erreur génération: {e}")  # Debug dans les logs
-        import traceback
-        traceback.print_exc()
-        return None
-
-# Interface Gradio
-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é")
+    pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    pipe.to("cuda")
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            lora_input = gr.Textbox(
-                label="Repo HuggingFace LoRA", 
-                placeholder="ex: stabilityai/stable-diffusion-xl-base-1.0",
-                value="XLabs-AI/flux-lora-collection"
-            )
-            subfolder_input = gr.Dropdown(
-                label="Style LoRA", 
-                choices=["realism", "anime", "scenery", "art", "disney"],
-                value="realism"
-            )
-            load_btn = gr.Button("Charger LoRA", variant="secondary")
-            status = gr.Textbox(label="Status", interactive=False)
-        
-        with gr.Column(scale=4):
-            prompt = gr.Textbox(
-                label="Prompt", 
-                placeholder="une belle image artistique, haute qualité, détaillée",
-                lines=3,
-                value="beautiful artistic portrait, high quality, detailed"
-            )
-            neg_prompt = gr.Textbox(
-                label="Prompt négatif", 
-                value="blurry, deformed, ugly, lowres, text, watermark"
-            )
+    if lora_repo and lora_path:
+        pipe.load_lora_weights(lora_path)
+        pipe.fuse_lora(lora_scale=lora_scale)  # Après to("cuda")
     
-    with gr.Row():
-        steps = gr.Slider(1, 20, value=4, label="Steps", step=1)
-        width = gr.Slider(512, 2048, value=1024, step=128, label="Largeur")
-        height = gr.Slider(512, 2048, value=1024, step=128, label="Hauteur")
-        lora_scale_slider = gr.Slider(0, 2, value=0.8, step=0.1, label="Force LoRA")
-        seed = gr.Number(value=-1, label="Seed (-1=random)", precision=0)
+    pipe.enable_model_cpu_offload()  # Optim VRAM
     
-    generate_btn = gr.Button("🚀 Générer Image", variant="primary", size="lg")
-    output = gr.Image(label="Résultat", type="pil")
+    generator = torch.Generator("cuda").manual_seed(seed if seed != -1 else random.randint(0, 2**32))
     
-    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], 
-        outputs=output
-    )
+    image = pipe(
+        prompt, 
+        negative_prompt=negative_prompt,
+        height=height, width=width, 
+        num_inference_steps=steps,
+        guidance_scale=0.0, 
+        generator=generator,
+        max_sequence_length=256  # Pour FLUX
+    ).images[0]
+    
+    del pipe  # Nettoyage
+    torch.cuda.empty_cache()
+    return image
+
+# Interface Gradio (inchangée)
+with gr.Blocks(title="Flux Schnell + LoRA") as demo:
+    # ... (votre interface reste identique)
+    pass
 
-if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+demo.launch()