Aggiunto nuovo appv3

app.py

@@ -3,7 +3,7 @@ import torch
 from PIL import Image
 from diffusers import StableDiffusionImg2ImgPipeline
 from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
+from safetensors import safe_open
 
 # Scarica il modello .safetensors dalla tua model card
 model_path = hf_hub_download(
@@ -11,17 +11,49 @@ model_path = hf_hub_download(
     filename="juggernaut_reborn.safetensors"
 )
 
-# Carica il modello Stable Diffusion di base
+# Carica il modello Stable Diffusion
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",  # Modello base compatibile
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     use_safetensors=True,
-    variant="fp16" if torch.cuda.is_available() else None,
+    safety_checker=None
 )
 
-# Carica i pesi dal file .safetensors usando la libreria safetensors
-state_dict = load_file(model_path)
-pipe.load_state_dict(state_dict, strict=False)
+# Determina se è un LoRA o un modello completo
+try:
+    # Opzione 1: Carica come LoRA
+    pipe.load_lora_weights(model_path)
+except Exception as lora_error:
+    try:
+        # Opzione 2: Carica direttamente nell'UNet
+        with safe_open(model_path, framework="pt", device="cpu") as f:
+            unet_params = {k: f.get_tensor(k) for k in f.keys() if k.startswith("unet.")}
+            
+        # Rimuovi il prefisso "unet." dalle chiavi se presente
+        clean_unet_params = {}
+        for k, v in unet_params.items():
+            if k.startswith("unet."):
+                clean_unet_params[k[5:]] = v  # Rimuovi il prefisso "unet."
+            else:
+                clean_unet_params[k] = v
+                
+        # Carica i pesi nell'UNet
+        if clean_unet_params:
+            pipe.unet.load_state_dict(clean_unet_params, strict=False)
+        else:
+            # Opzione 3: Prova a caricare come checkpoint completo
+            pipe = StableDiffusionImg2ImgPipeline.from_single_file(
+                model_path,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                use_safetensors=True,
+                load_safety_checker=False
+            )
+    except Exception as e:
+        print(f"Errore nel caricamento del modello: {e}")
+        # Fallback: usa il modello base
+        print("Utilizzo del modello base come fallback")
+
+# Sposta il modello su GPU se disponibile
 pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
 def enhance_image(image, prompt, negative_prompt, cfg_scale, denoising_strength):