Update app.py

app.py

@@ -1,37 +1,34 @@
-# Function to generate an image from a text prompt
-def generate_image_from_prompt(prompt, generator, device, z_dim):
-    # A very simple way to map a prompt to a latent vector.
-    # In a real application, you would use a more sophisticated text encoder.
-    # Here, we just use a hash of the prompt to seed a random number generator.
-    import hashlib
-    seed = int(hashlib.sha256(prompt.encode('utf-8')).hexdigest(), 16) % (2**32 - 1)
-    torch.manual_seed(seed)
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import gradio as gr
+from model import Generator
 
-    noise = torch.randn(1, z_dim, 1, 1).to(device)
-    with torch.no_grad():
-        fake_img = generator(noise).detach().cpu().squeeze(0)
-        img = (fake_img.permute(1, 2, 0) + 1) / 2
-    return img
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Configurações
+latent_dim = 100
+generator = Generator(latent_dim).to(device)
+generator.load_state_dict(torch.load("generator128.pth", map_location=device))
+generator.eval()
 
-# Load the trained generator model
-# Make sure to run the training cell (VDcHP883uCpZ) at least once to save the model
-gen_file = os.path.join(model_path, "generator128.pth")
-if os.path.exists(gen_file):
-    # Initialize the generator model
-    G = Generator().to(device)
-    G.load_state_dict(torch.load(gen_file))
-    G.eval() # Set the generator to evaluation mode
-    print("✅ Generator model loaded successfully!")
-else:
-    print("⚠️ Generator model not found. Please run the training cell first.")
+# Função para gerar imagem
+def generate_image(seed: int):
+    torch.manual_seed(seed)
+    noise = torch.randn(1, latent_dim, 1, 1, device=device)
+    with torch.no_grad():
+        fake_image = generator(noise).detach().cpu()
+    fake_image = (fake_image + 1) / 2  # Traz para [0, 1]
+    to_pil = transforms.ToPILImage()
+    return to_pil(fake_image.squeeze(0))
 
-# Example usage:
-if os.path.exists(gen_file):
-    prompt = "cat and dog"
-    generated_image = generate_image_from_prompt(prompt, G, device, z_dim)
+# Interface Gradio
+demo = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Slider(0, 99999, value=42, label="Seed"),
+    outputs=gr.Image(type="pil", label="Generated Image"),
+    title="ZYI 0.1 - CringeCoin Generator",
+    description="Enter a seed to generate a unique 128x128 image with ZYI 0.1."
+)
 
-    plt.figure(figsize=(4, 4))
-    plt.axis("off")
-    plt.title(f"Generated Image for Prompt: '{prompt}'")
-    plt.imshow(generated_image)
-    plt.show()
+demo.launch()