Spaces:
Running on Zero
Running on Zero
| # ============================================================================== | |
| # 1. DEPENDANCES & COMPATIBILITE COLAB / ZERO-GPU | |
| # ============================================================================== | |
| #!pip install -q diffusers transformers accelerate datasets gradio torch torchvision huggingface_hub | |
| try: | |
| import spaces | |
| except ImportError: | |
| class MockSpaces: | |
| def GPU(self, fn=None, duration=None): return fn if fn is not None else lambda f: f | |
| spaces = MockSpaces() | |
| import os, json, torch, torch.nn as nn, numpy as np, gradio as ui | |
| from PIL import Image | |
| from datasets import load_dataset | |
| from torchvision import transforms | |
| from transformers import CLIPTextModel, CLIPTokenizer | |
| from diffusers import UNet2DConditionModel, AutoencoderKL, DDIMScheduler | |
| from huggingface_hub import HfApi | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE = None, None, None, None, None, 64 | |
| # ============================================================================== | |
| # 2. INITIALISATION ARCHITECTURE LOU DE ZERO | |
| # ============================================================================== | |
| def init_lou_architecture(image_size): | |
| tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") | |
| text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") | |
| vae = AutoencoderKL( | |
| in_channels=3, out_channels=3, latent_channels=4, | |
| down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D"), | |
| up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D"), block_out_channels=(64, 128) | |
| ) | |
| unet = UNet2DConditionModel( | |
| sample_size=image_size // 8, in_channels=4, out_channels=4, layers_per_block=1, | |
| block_out_channels=(64, 128), down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), | |
| up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), cross_attention_dim=768 | |
| ) | |
| noise_scheduler = DDIMScheduler(num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear") | |
| return tokenizer, text_encoder, unet, vae, noise_scheduler | |
| # ============================================================================== | |
| # 3. TRAIN ET EXPORT REPO STRUCTURÉ (VAE + UNET COMPLET) | |
| # ============================================================================== | |
| def train_lou_model(repo, token, dataset_path, lr, epochs, size, progress=ui.Progress()): | |
| global LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE | |
| LOU_IMAGE_SIZE = int(size) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_UNET, LOU_VAE, LOU_SCHEDULER = init_lou_architecture(LOU_IMAGE_SIZE) | |
| LOU_UNET.to(device); LOU_TEXT_ENCODER.to(device); LOU_VAE.to(device) | |
| dataset = load_dataset(dataset_path, split="train") | |
| tf = transforms.Compose([transforms.Resize((LOU_IMAGE_SIZE, LOU_IMAGE_SIZE)), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]) | |
| optimizer = torch.optim.AdamW(list(LOU_UNET.parameters()) + list(LOU_VAE.parameters()), lr=float(lr)) | |
| LOU_UNET.train(); LOU_VAE.train() | |
| for epoch in range(int(epochs)): | |
| for batch in progress.tqdm(dataset, desc=f"Epoch {epoch + 1}/{epochs}"): | |
| optimizer.zero_grad() | |
| clean = tf(batch["image"].convert("RGB")).unsqueeze(0).to(device) | |
| latents = LOU_VAE.encode(clean).latent_dist.sample() * 0.18215 | |
| noise = torch.randn_like(latents) | |
| t = torch.randint(0, 1000, (1,), device=device).long() | |
| noisy = LOU_SCHEDULER.add_noise(latents, noise, t) | |
| txt = LOU_TOKENIZER(batch["text"], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device) | |
| loss = nn.functional.mse_loss(LOU_UNET(noisy, t, LOU_TEXT_ENCODER(txt).last_hidden_state).sample, noise) | |
| loss.backward(); optimizer.step() | |
| base_dir = "./lou_model_repo" | |
| os.makedirs(base_dir, exist_ok=True) | |
| LOU_UNET.save_pretrained(os.path.join(base_dir, "unet")) | |
| LOU_VAE.save_pretrained(os.path.join(base_dir, "vae")) | |
| LOU_TEXT_ENCODER.save_pretrained(os.path.join(base_dir, "text_encoder")) | |
| LOU_TOKENIZER.save_pretrained(os.path.join(base_dir, "tokenizer")) | |
| LOU_SCHEDULER.save_pretrained(os.path.join(base_dir, "scheduler")) | |
| with open(os.path.join(base_dir, "model_index.json"), "w") as f: | |
| json.dump({"_class_name": "LouDiffusionArchitecture", "_diffusers_version": "0.21.4", "scheduler": ["diffusers", "DDIMScheduler"], "text_encoder": ["transformers", "CLIPTextModel"], "tokenizer": ["transformers", "CLIPTokenizer"], "unet": ["diffusers", "UNet2DConditionModel"], "vae": ["diffusers", "AutoencoderKL"]}, f, indent=2) | |
| if token and repo: | |
| try: | |
| api = HfApi(); api.create_repo(repo_id=repo, token=token, exist_ok=True) | |
| api.upload_folder(folder_path=base_dir, repo_id=repo, repo_type="model", token=token) | |
| return f"✅ Entraîné et synchronisé sur HF via hfapi." | |
| except Exception as e: return f"⚠️ Structure locale ok, erreur hfapi : {str(e)}" | |
| return f"✅ Entraîné ! Structure locale créée dans : {base_dir}" | |
| # ============================================================================== | |
| # 4. INFERENCE EN ESPACE LATENT (VITESSE MAXIMALE FP16) | |
| # ============================================================================== | |
| def generate_lou_images(prompt, steps, num_images, seed, cfg, auto): | |
| global LOU_UNET, LOU_VAE, LOU_TOKENIZER, LOU_TEXT_ENCODER, LOU_SCHEDULER, LOU_IMAGE_SIZE | |
| if LOU_UNET is None: return [Image.new("RGB", (64, 64), color="gray")], "❌ Lancez d'abord l'entraînement." | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| LOU_UNET.to(device=device, dtype=torch.float16); LOU_TEXT_ENCODER.to(device=device, dtype=torch.float16); LOU_VAE.to(device=device, dtype=torch.float16) | |
| LOU_UNET.eval(); LOU_TEXT_ENCODER.eval(); LOU_VAE.eval() | |
| if seed >= 0: torch.manual_seed(int(seed)); np.random.seed(int(seed)) | |
| t_in = LOU_TOKENIZER([prompt], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device) | |
| u_in = LOU_TOKENIZER([""], padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device) | |
| with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16): | |
| context = torch.cat([LOU_TEXT_ENCODER(u_in).last_hidden_state, LOU_TEXT_ENCODER(t_in).last_hidden_state]) | |
| generated = [] | |
| LOU_SCHEDULER.set_timesteps(int(steps)) | |
| with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16): | |
| for _ in range(int(num_images)): | |
| latents = torch.randn((1, 4, LOU_IMAGE_SIZE // 8, LOU_IMAGE_SIZE // 8), device=device, dtype=torch.float16) | |
| for t in LOU_SCHEDULER.timesteps: | |
| p_uncond, p_text = LOU_UNET(torch.cat([latents] * 2), t, encoder_hidden_states=context).sample.chunk(2) | |
| latents = LOU_SCHEDULER.step(p_uncond + cfg * (p_text - p_uncond), t, latents).prev_sample | |
| img_tensor = (LOU_VAE.decode(latents / 0.18215).sample / 2 + 0.5).clamp(0, 1).squeeze(0) | |
| generated.append(Image.fromarray((img_tensor.cpu().float().permute(1, 2, 0).numpy() * 255).astype(np.uint8))) | |
| return generated, f"🎉 Images {LOU_IMAGE_SIZE}x{LOU_IMAGE_SIZE} générées par Lou !" | |
| # ============================================================================== | |
| # 5. CODE GRADIO INTERFACE MINIMALISTE COMPACTE | |
| # ============================================================================== | |
| with ui.Blocks(theme=ui.themes.Default()) as app: | |
| with ui.Tab("🏋️ Train"): | |
| with ui.Row(): | |
| repo_i = ui.Textbox(label="Repo HF", placeholder="user/repo"); token_i = ui.Textbox(label="Token HF", type="password") | |
| ds_i = ui.Textbox(label="Dataset", value="diffusers/pokemon-gpt4-captions"); lr_i = ui.Number(label="LR", value=1e-4) | |
| ep_i = ui.Slider(1, 50, value=3, step=1, label="Époques"); sz_i = ui.Slider(32, 1024, value=64, step=32, label="Taille Image") | |
| t_btn = ui.Button("🚀 Lancer l'entraînement de Lou", variant="primary"); t_out = ui.Textbox(label="Statut", lines=1) | |
| with ui.Tab("🎨 Playground"): | |
| with ui.Row(): | |
| pr_i = ui.Textbox(label="Prompt", value="a cute pokemon", scale=2); st_i = ui.Slider(5, 100, value=20, step=1, label="Steps") | |
| num_i = ui.Slider(1, 4, value=1, step=1, label="Images"); sd_i = ui.Number(label="Seed", value=-1) | |
| gd_i = ui.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance"); au_i = ui.Checkbox(label="Autotraining", value=True) | |
| g_btn = ui.Button("🔮 Générer", variant="secondary"); out_gal = ui.Gallery(label="Outputs", height=200); g_out = ui.Textbox(label="Statut Gen", lines=1) | |
| t_btn.click(train_lou_model, inputs=[repo_i, token_i, ds_i, lr_i, ep_i, sz_i], outputs=t_out) | |
| g_btn.click(generate_lou_images, inputs=[pr_i, st_i, num_i, sd_i, gd_i, au_i], outputs=[out_gal, g_out]) | |
| app.queue(default_concurrency_limit=4) | |
| app.launch(inline=True, share=True,debug=True) |