add gen_image

gen_image.py

@@ -0,0 +1,34 @@
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
+import torch
+
+class ImageGenerator:
+    def __init__(self, model_id="stabilityai/stable-diffusion-2-1-base", device="cuda"):
+        scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            model_id, 
+            scheduler=scheduler, 
+            torch_dtype=torch.float16
+        )
+        self.pipe = self.pipe.to(device)
+    
+    def generate(self, prompt, negative_prompt=None, output_path=None):
+        image = self.pipe(prompt, negative_prompt=negative_prompt).images[0]
+        
+        if output_path:
+            image.save(output_path)
+            
+        return image
+
+
+# Example usage
+if __name__ == "__main__":
+    generator = ImageGenerator()
+    import time
+    start_time = time.time()
+    image = generator.generate(
+        prompt="magenta trapezoids layered on a transluscent silver sheet, simple, icon",
+        negative_prompt="3d, blurry, complex geometry, realistic",
+        output_path="sheet.png"
+    )
+    end_time = time.time()
+    print(f"Time taken: {end_time - start_time} seconds")

requirements.txt

@@ -16,6 +16,9 @@ accelerate
 openai
 tqdm
 dotenv
+diffusers
+safetensors
+xformers
 
 # pip install 'tensorflow[and-cuda]'
 # pip install git+https://github.com/openai/CLIP.git