File size: 5,350 Bytes
38b4eff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import os
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from PIL import Image
import time
import requests
from io import BytesIO
class NeuralAIDiffusion:
def __init__(self, model_id="runwayml/stable-diffusion-v1-5", device=None):
self.model_id = model_id
if device:
self.device = device
else:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.pipe = None
self.img2img_pipe = None
self.is_loaded = False
print(f"[NeuralAI Diffusion] Initialized on {self.device}")
def load_model(self, mode="text2img"):
if self.is_loaded and (self.pipe if mode == "text2img" else self.img2img_pipe):
return
print(f"[NeuralAI Diffusion] Loading {mode} model {self.model_id}...")
try:
# Using float32 for CPU to avoid errors, float16 for CUDA
dtype = torch.float16 if self.device == "cuda" else torch.float32
if mode == "text2img":
self.pipe = StableDiffusionPipeline.from_pretrained(
self.model_id,
torch_dtype=dtype,
safety_checker=None # Disable safety checker for faster loading if needed, or keep for safety
)
self.pipe.to(self.device)
# Optimization for CPU
if self.device == "cpu":
self.pipe.enable_attention_slicing()
else:
self.img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
self.model_id,
torch_dtype=dtype,
safety_checker=None
)
self.img2img_pipe.to(self.device)
if self.device == "cpu":
self.img2img_pipe.enable_attention_slicing()
self.is_loaded = True
print(f"[NeuralAI Diffusion] {mode} model loaded successfully.")
except Exception as e:
print(f"[NeuralAI Diffusion] Error loading model: {e}")
if self.model_id != "segmind/tiny-sd":
print("[NeuralAI Diffusion] Attempting fallback to tiny-sd...")
self.model_id = "segmind/tiny-sd"
self.load_model(mode)
def generate(self, prompt, output_path, negative_prompt=None, num_steps=20, guidance_scale=7.5):
self.load_model("text2img")
# Enhanced Prompting for "Better Images"
quality_boost = "masterpiece, high quality, 8k, highly detailed, professional photography"
if "moon" in prompt.lower():
quality_boost += ", sharp craters, lunar surface detail, space background, realistic"
full_prompt = f"{prompt}, {quality_boost}"
if negative_prompt is None:
negative_prompt = "blurry, low quality, distorted, watermark, text, grainy, low resolution"
print(f"[NeuralAI Diffusion] Generating: {full_prompt}")
start_time = time.time()
try:
image = self.pipe(
prompt=full_prompt,
negative_prompt=negative_prompt,
num_inference_steps=num_steps,
guidance_scale=guidance_scale
).images[0]
image.save(output_path)
print(f"[NeuralAI Diffusion] Image saved to {output_path} (took {time.time() - start_time:.2f}s)")
return True
except Exception as e:
print(f"[NeuralAI Diffusion] Generation failed: {e}")
return False
def transform(self, prompt, image_path, output_path, strength=0.75, num_steps=20):
self.load_model("img2img")
quality_boost = "masterpiece, high quality, 8k, highly detailed"
full_prompt = f"{prompt}, {quality_boost}"
print(f"[NeuralAI Diffusion] Transforming image with prompt: {full_prompt}")
start_time = time.time()
try:
if image_path.startswith("http"):
response = requests.get(image_path)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
else:
init_image = Image.open(image_path).convert("RGB")
init_image = init_image.resize((512, 512))
image = self.img2img_pipe(
prompt=full_prompt,
image=init_image,
strength=strength,
num_inference_steps=num_steps
).images[0]
image.save(output_path)
print(f"[NeuralAI Diffusion] Transformed image saved to {output_path} (took {time.time() - start_time:.2f}s)")
return True
except Exception as e:
print(f"[NeuralAI Diffusion] Transformation failed: {e}")
return False
if __name__ == "__main__":
import sys
mode = sys.argv[1] if len(sys.argv) > 1 else "gen"
prompt = sys.argv[2] if len(sys.argv) > 2 else "A high-tech AI logo"
output = sys.argv[3] if len(sys.argv) > 3 else "output.png"
engine = NeuralAIDiffusion()
if mode == "edit" and len(sys.argv) > 4:
engine.transform(prompt, sys.argv[4], output)
else:
engine.generate(prompt, output)
|