Update README.md

README.md

@@ -52,53 +52,52 @@ While the current version (v1) produces stylistic, slightly "painterly" or "pixe
 ## 🛠️ How to use
 ```python
 import torch
-from transformers import AutoTokenizer, CLIPTextModel, AutoConfig
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
 import matplotlib.pyplot as plt
 import numpy as np
-
-# import classes
+import os
+import shutil
+from transformers import AutoTokenizer, CLIPTextModel, AutoModel, AutoConfig
 
 def generate_fixed_from_hub(prompt, model_id="TopAI-1/Pixel-1"):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"🚀 Working on {device}...")
 
-    # 1. Download
-    print("📥 Downloading weights directly from Hub...")
-    weights_path = hf_hub_download(repo_id=model_id, filename="model.safetensors")
-    config_path = hf_hub_download(repo_id=model_id, filename="config.json")
+    # 1. ניקוי Cache כדי לוודא שאתה מושך את התיקונים החדשים מה-Hub
+    cache_path = os.path.expanduser(f"~/.cache/huggingface/hub/models--{model_id.replace('/', '--')}")
+    if os.path.exists(cache_path):
+        print("🧹 Clearing old cache to fetch your latest fixes...")
+        shutil.rmtree(cache_path)
 
-    # 2. loading CLIP
+    # 2. טעינת CLIP
     tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14")
     text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
 
-    # 3. load wieghts
-    config = TopAIImageConfig.from_pretrained(model_id)
-    model = TopAIImageGenerator(config)
-    
-    state_dict = load_file(weights_path, device=device)
+    # 3. טעינת המודל והקונפיג אוטומטית מה-Hub
+    # בזכות ה-auto_map ב-config.json, transformers ימצא לבד את המחלקות
+    print("📥 Downloading architecture and weights directly from Hub...")
+    model = AutoModel.from_pretrained(
+        model_id, 
+        trust_remote_code=True, 
+        force_download=True
+    ).to(device)
     
-    clean_sd = {k.replace('\xa0', ' '): v for k, v in state_dict.items()}
-    
-    # Loading
-    model.load_state_dict(clean_sd, strict=False) 
-    model.to(device).eval()
-    print("✅ Weights loaded perfectly!")
+    model.eval()
+    print("✅ Model loaded successfully!")
 
-    # 4. Generation
+    # 4. יצירה
+    print(f"🎨 Generating: {prompt}")
     inputs = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to(device)
+    
     with torch.no_grad():
         emb = text_encoder(inputs.input_ids).pooler_output
         out = model(emb)
 
-    # 5. Show
+    # 5. תצוגה
     img = (out.squeeze(0).cpu().permute(1, 2, 0).numpy() + 1.0) / 2.0
     plt.figure(figsize=(8, 8))
     plt.imshow(np.clip(img, 0, 1))
     plt.axis('off')
-    plt.title(prompt)
     plt.show()
 
-# Run
+# הרצה
 generate_fixed_from_hub("Window with metal bars and fence shadow")