Initial commit

app.py

@@ -2,40 +2,35 @@ import gradio as gr
 import torch
 import numpy as np
 import random
-from diffusers import FluxImg2ImgPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
 
-dtype = torch.bfloat16
+dtype = torch.float16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 
-MODEL_ID = "black-forest-labs/FLUX.1-schnell"
-LORA_REPO = "alvdansen/flux-koda"
-LORA_FILE = "flux_dev_koda_araminta_k.safetensors"
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
 
 STYLES = {
     "Anime": {
-        "lora": "alvdansen/flux-koda",
-        "file": "flux_dev_koda_araminta_k.safetensors",
-        "prompt": "anime style portrait, detailed face, vibrant colors, high quality"
+        "prompt": "anime style portrait, detailed face, vibrant colors, high quality, studio ghibli"
     },
     "Pixel Art": {
-        "lora": "nerijs/pixel-art-xl",
-        "file": "pixel-art-xl.safetensors",
-        "prompt": "pixel art avatar, 16-bit style, retro game character"
+        "prompt": "pixel art avatar, 16-bit style, retro game character, sharp pixels"
     },
     "Ghibli": {
-        "lora": "aleksa-codes/flux-ghibsky-illustration",
-        "file": "flux_ghibsky_illustration.safetensors",
-        "prompt": "ghibli studio style illustration, soft colors, dreamy atmosphere"
+        "prompt": "ghibli studio style illustration, soft colors, dreamy atmosphere, watercolor"
     },
 }
 
 print("Chargement du modèle...")
-pipe = FluxImg2ImgPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype)
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=dtype,
+    safety_checker=None
+)
 pipe.to(device)
-
-current_style = None
+print("Modèle chargé !")
 
 def generate_avatar(
     input_image,
@@ -45,33 +40,23 @@ def generate_avatar(
     randomize_seed,
     progress=gr.Progress(track_tqdm=True)
 ):
-    global current_style
-
     if input_image is None:
         raise gr.Error("Veuillez uploader une photo !")
 
-    # Charger le bon LoRA selon le style
-    if style_name != current_style:
-        style = STYLES[style_name]
-        pipe.unload_lora_weights()
-        pipe.load_lora_weights(style["lora"], weight_name=style["file"])
-        current_style = style_name
-
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator(device=device).manual_seed(seed)
     prompt = STYLES[style_name]["prompt"]
 
-    # Redimensionner l'image input
     input_image = input_image.resize((512, 512))
 
     result = pipe(
         prompt=prompt,
         image=input_image,
-        strength=strength,       # 0.5 = garde ressemblance, 0.9 = full style
+        strength=strength,
         num_inference_steps=20,
-        guidance_scale=3.5,
+        guidance_scale=7.5,
         generator=generator,
     ).images[0]