Update README.md

README.md

@@ -51,13 +51,54 @@ While the current version (v1) produces stylistic, slightly "painterly" or "pixe
 
 ## 🛠️ How to use
 ```python
-from transformers import AutoTokenizer, CLIPTextModel
+import torch
+from transformers import AutoTokenizer, CLIPTextModel, AutoConfig
+from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-
-# Load the custom Pixel-1 Generator
-model = TopAIImageGenerator.from_pretrained("Your-Repo/Pixel-1")
-tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14")
-
-# Generate
-prompt = "Your creative prompt here"
-# ... (standard inference code)
+import matplotlib.pyplot as plt
+import numpy as np
+
+# import classes
+
+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")
+
+    # 2. loading 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)
+    
+    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!")
+
+    # 4. Generation
+    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
+    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")