Create hk.py

hk.py

@@ -0,0 +1,61 @@
+import os
+import requests
+from stable_diffusion_cpp import StableDiffusion
+
+# Function to download a file
+def download_file(url, save_path):
+    response = requests.get(url)
+    response.raise_for_status()  # Check for request errors
+    with open(save_path, 'wb') as f:
+        f.write(response.content)
+
+# Define model URLs
+vae_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/vae/diffusion_pytorch_model.safetensors"
+model_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/flowgram.safetensors"
+clip_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/text_encoder/clip_l.safetensors"
+t5xxl_url = "https://huggingface.co/pranavajay/flowgram/resolve/main/text_encoder/t5xxl_fp16.safetensors"
+
+# Define local save paths
+vae_path = "vae/diffusion_pytorch_model.safetensors"
+model_path = "flowgram.safetensors"
+clip_path = "text_encoder/clip_l.safetensors"
+t5xxl_path = "text_encoder/t5xxl_fp16.safetensors"
+
+# Create directories if they don't exist
+os.makedirs("vae", exist_ok=True)
+os.makedirs("text_encoder", exist_ok=True)
+
+# Download the models if they do not exist
+if not os.path.exists(vae_path):
+    download_file(vae_url, vae_path)
+if not os.path.exists(model_path):
+    download_file(model_url, model_path)
+if not os.path.exists(clip_path):
+    download_file(clip_url, clip_path)
+if not os.path.exists(t5xxl_path):
+    download_file(t5xxl_url, t5xxl_path)
+
+# Initialize the StableDiffusion model
+flowgram_diffusion = StableDiffusion(
+    diffusion_model_path=model_path,
+    clip_l_path=clip_path,
+    t5xxl_path=t5xxl_path,
+    vae_path=vae_path,
+)
+
+# Function to generate an image from text
+def generate_image(prompt, num_images=1, guidance_scale=7.5):
+    # Generate images
+    images = flowgram_diffusion.generate(prompt, num_images=num_images, guidance_scale=guidance_scale)
+    
+    # Return the generated images
+    return images
+
+# Example usage
+if __name__ == "__main__":
+    prompt = "A beautiful landscape with mountains and a river"
+    generated_images = generate_image(prompt)
+    
+    # Save or display the images
+    for i, img in enumerate(generated_images):
+        img.save(f"generated_image_{i}.png")  # Save each generated image