raySD 0.1.0__py3-none-any.whl

raySD/__init__.py

@@ -0,0 +1,6 @@
+from .pipeline import *
+from .prompt import *
+from .pydantic_model import *
+from .utils import *
+
+__version__ = "0.0.1"

raySD/pipeline/Abstract_Pipeline.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+
+class AbstractPipeline(ABC):
+    def __init__(self, **kwargs):
+         ## Hook model
+        self.module_pre_hook = []
+        self.module_post_hook = []
+
+    def __call__(self, *args, **kwargs):
+        for hook in self.module_pre_hook:
+            x = hook(*args, **kwargs)
+        x = self.forward(*args, **kwargs)
+        for hook in self.module_post_hook:
+            x = hook(*args, **kwargs)
+        return x
+    
+    def register_pre_hook(self, hook):
+        self.module_pre_hook.append(hook)
+
+    def register_post_hook(self, hook):
+        self.module_post_hook.append(hook)
+
+    def forward(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def load_weights(self):
+        pass

raySD/pipeline/Stable_Diffusion_IP_Adapter.py

@@ -0,0 +1,115 @@
+import numpy as np
+import torch
+from raySD.utils.Enable_xFormers import *
+from raySD.pipeline.Abstract_Pipeline import *
+from raySD.prompt.chibi_prompt import *
+from raySD.pydantic_model.Stable_Diffusion_IP_Adapter_Model import *
+import time
+import cv2
+from PIL import Image
+import importlib
+
+class StableDiffusionIPAdapterPipeline(AbstractPipeline):
+    def __init__(self, **kwargs):
+        super().__init__()
+        ## Config
+        self.config = SD_IP_Config(**kwargs)
+
+    def update_params_infer(self, **kwargs):
+        self.config.inference_steps = kwargs.get("inference_steps", 20)
+        self.config.guidance_scale = kwargs.get("guidance_scale", 7)
+        self.config.ip_adapter_scale = kwargs.get("ip_adapter_scale", 1)
+        self.config.clip_skip=kwargs.get("clip_skip", 2)
+        self.config.num_images_per_prompt=kwargs.get("num_images_per_prompt", 1)
+
+    def load_weights(self):
+        from diffusers import DPMSolverMultistepScheduler
+        from diffusers import StableDiffusionPipeline
+        
+        ## Device
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"DEVICE: {self.device}")
+        ## Load Stable Diffusion model
+        pipeline = StableDiffusionPipeline.from_single_file(
+            self.config.checkpoint_path,
+            torch_dtype=torch.float16,
+            safety_checker=None
+        ).to(self.device)
+
+        ### Scheduler
+        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipeline.scheduler.config,
+            algorithm_type="dpmsolver++",       # <-- DPM++
+            use_karras_sigmas=True,             # <-- Karras
+            solver_order=2                      # <-- 2M (2nd order multistep)
+        )
+
+        if self.config.lora_style_path:
+            print("Load LoRA done")
+            pipeline.load_lora_weights(self.config.lora_style_path)
+
+        enable_xformers_if_available(pipeline)
+        
+        ### Load pipeline to IPAdapter
+        if self.config.ip_adapter_plus:
+            from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlus
+            self.ip_adapter = IPAdapterFaceIDPlus(
+                pipeline,
+                image_encoder_path=self.config.image_encoder_path,
+                ip_ckpt=self.config.ip_adapter_checkpoint_path,
+                device=self.device
+            )   
+        else:
+            from ip_adapter.ip_adapter_faceid import IPAdapterFaceID
+            self.ip_adapter = IPAdapterFaceID(
+                pipeline,
+                ip_ckpt=self.config.ip_adapter_checkpoint_path,
+                device=self.device
+            )
+
+        module = importlib.import_module("insightface.app")
+        # Load Face Analysis
+        self.face_app = module.FaceAnalysis(name=self.config.insightface_app_name, providers=self.config.insightface_providers)
+        self.face_app.prepare(ctx_id=0, det_size=self.config.det_size_face)
+        
+        module_util = importlib.import_module("insightface.utils")
+        self.face_align = getattr(module_util, "face_align")
+
+        print("LOAD SD PIPELINE DONE")
+    
+    def forward(self, x: Image.Image, **kwargs)->Image.Image:
+        img = cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR)
+        faces = self.face_app.get(img)
+        face_emb = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0).to(self.device)
+        aligned_face = self.face_align.norm_crop(img, landmark=faces[0].kps, image_size=self.config.target_size)
+
+        print(face_emb.shape)
+        print(aligned_face.shape)
+        
+        gender = kwargs["gender"]
+        age = kwargs["age"]
+
+        if self.config.style_name == "chibi_style":
+            chibi_prompt = ChibiPrompt2(gender=gender)
+            self.prompt, self.negative_prompt = chibi_prompt.get_prompt()
+        
+        start_time = time.time()
+        image = self.ip_adapter.generate(
+            face_image=aligned_face,
+            faceid_embeds=face_emb,
+            prompt=self.prompt,
+            negative_prompt=self.negative_prompt,
+            guidance_scale=self.config.guidance_scale,
+            num_inference_steps=self.config.inference_steps,
+            scale=self.config.ip_adapter_scale,
+            num_samples=self.config.num_images_per_prompt,
+            width=self.config.width,
+            height=self.config.height,
+        )[0]
+        end_time = time.time()
+        print(f"Inference time: {end_time - start_time:.2f} seconds")
+        return image
+
+
+
+

raySD/pipeline/Super_Resolution_Pipeline.py

@@ -0,0 +1,42 @@
+from raySD.pipeline.Abstract_Pipeline import *
+from PIL import Image
+import importlib
+import time
+import numpy as np
+import cv2
+
+class SuperResolutionPipeline(AbstractPipeline):
+    def __init__(self, model_path: str, device="CPU", output_size=(768, 768)):
+        super().__init__()
+        self.model_path = model_path
+        self.device = device
+        self.output_size = output_size
+        
+    def load_weights(self):
+        module = importlib.import_module("openvino")
+        Core = getattr(module, "Core")
+        core = Core()
+        model_ov = core.read_model(self.model_path)
+        self.model = core.compile_model(model_ov, device_name=self.device)
+        print("LOAD SUPER RESOLUTION MODEL DONE")
+
+    
+    def forward(self, x: Image.Image, **kwargs)->Image.Image:
+        st = time.time()
+        x = np.array(x)
+        x = x / 255.0
+        x = np.expand_dims(x, axis=0)
+        x = np.transpose(x, (0, 3, 1, 2))
+        y = self.model([x])
+        start_time = time.time()
+        y = y[next(iter(y))]
+        end_time = time.time()
+        print(f"Inference ESRGAN: {end_time - start_time}s")
+        y = np.transpose(y, (0, 2, 3, 1))
+        y = np.squeeze(y)
+        y = cv2.resize(y, self.output_size, interpolation=cv2.INTER_AREA)
+        y = cv2.convertScaleAbs(y*255)
+        image_out = Image.fromarray(y)
+        ed = time.time()
+        print(f"All time IF: {ed - st}s")
+        return image_out

raySD/pipeline/__init__.py

@@ -0,0 +1,2 @@
+from .Stable_Diffusion_IP_Adapter import *
+from .Super_Resolution_Pipeline import *

raySD/prompt/__init__.py

@@ -0,0 +1 @@
+from .chibi_prompt import *

raySD/prompt/chibi_prompt.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel, model_validator
+from typing import List, Dict, Any, Optional
+
+class ChibiPrompt():
+    def __init__(self, gender: int):
+        self.gender = gender
+        if gender == 1:
+            self.key1 = "boy"
+            self.key2 = "male"
+        else:
+            self.key1 = "girl"
+            self.key2 = "female"
+
+        # self.prompt = f"chibi style, solo, looking at viewer, simple background, brown hair, blue shirt, 1{self.key1}, {self.key2} focus, brown hair, warm and friendly expression, natural smile, well-aligned teeth, realistic lips, smooth facial features, expressive eyes, soft lighting, high-quality shading, professional rendering, best quality, high resolution, portrait, blue background, natural proportions, detailed face, gentle grin"
+        # self.negative_prompt = f"worst quality, low quality, normal quaworst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad smile, bad anatomylity, jpeg artifacts, signature, watermark, username, blurry, bad anatomy, blurry teeth, deformed teeth, extra teeth, missing teeth"
+
+        self.prompt = f"chibi style, solo, looking at viewer, simple background, brown hair, blue shirt, 1{self.key1},  {self.key2} (20-30 years old), {self.key2} focus, brown hair, warm and friendly expression, natural smile, well-aligned teeth, realistic lips, smooth facial features, expressive eyes, soft lighting, high-quality shading, professional rendering, best quality, high resolution, portrait, blue background, natural proportions, detailed face, gentle grin"
+        self.negative_prompt = f"worst quality, low quality, normal quaworst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad smile, bad anatomylity, jpeg artifacts, signature, watermark, username, blurry, bad anatomy, blurry teeth, deformed teeth, extra teeth, missing teeth"
+
+    def get_prompt(self):
+        return self.prompt, self.negative_prompt
+    
+
+class ChibiPrompt1():
+    def __init__(self, gender: int):
+        self.gender = gender
+        if gender == 1:
+            self.key1 = "boy"
+            self.key2 = "male"
+        else:
+            self.key1 = "girl"
+            self.key2 = "female"
+
+        self.prompt = f"chibi style, solo, 1{self.key1}, blue shirt, brown hair, looking at viewer, warm and friendly expression, natural smile, realistic lips, well-aligned teeth, smooth facial features, expressive eyes, soft lighting, detailed face, gentle grin, portrait, blue background, simple background, professional rendering, high-quality shading, bright face, soft glow, best quality, high resolution"
+        self.negative_prompt = f"worst quality, low quality, normal quaworst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad smile, bad anatomylity, jpeg artifacts, signature, watermark, username, blurry, bad anatomy, blurry teeth, deformed teeth, extra teeth, missing teeth"
+
+    def get_prompt(self):
+        return self.prompt, self.negative_prompt
+
+
+class ChibiPrompt2():
+    def __init__(self, gender: int):
+        self.gender = gender
+        if gender == 1:
+            self.key1 = "boy"
+            self.key2 = "male"
+        else:
+            self.key1 = "girl"
+            self.key2 = "female"
+
+        self.prompt = f"chibi style, solo, 1{self.key1}, looking at viewer, blue shirt, friendly expression, natural smile, smooth fair skin, realistic lips, big round eyes, bright eyes, evenly lit bright face, soft glow, ambient diffuse light, high quality shading, detailed soft face, professional render, best quality, high resolution, blue background"
+
+        self.negative_prompt = f"worst quality, low quality, jpeg artifacts, signature, watermark, username, blurry, bad smile, bad anatomy, blurry teeth, deformed teeth, extra teeth, missing teeth, uneven lighting, harsh shadow, flat shading, discolored skin, dark chin, underexposed face, unnatural skin tone, narrow eyes, slit eyes, squinting eyes, closed eyes, color bleeding, blue tint on face"
+
+    def get_prompt(self):
+        return self.prompt, self.negative_prompt

raySD/pydantic_model/Image_Generate_Model.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel, model_validator
+from typing import List, Dict, Any, Optional
+
+class ImageGenerateModel(BaseModel):
+    image: str = ""
+    gender: int = -1
+    age: int = 0

raySD/pydantic_model/Image_Response_Model.py

@@ -0,0 +1,4 @@
+from pydantic import BaseModel, Field
+
+class ImageResponseModel(BaseModel):
+    image: str = ""

raySD/pydantic_model/Message_Model.py

@@ -0,0 +1,4 @@
+from pydantic import BaseModel, Field
+
+class MessageModel(BaseModel):
+    message: str = ""

raySD/pydantic_model/Stable_Diffusion_IP_Adapter_Model.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel, model_validator
+from typing import List, Dict, Any, Optional
+
+class SD_IP_Config(BaseModel):
+    pretrain_model_path: str = ""
+    checkpoint_path: str = ""
+    lora_style_path: str = ""
+    style_name: str = ""
+    height: int = 512
+    width: int = 512
+    ip_adapter_plus: bool = False
+    version: float = 1.0
+    image_encoder_path: str = ""
+    ip_adapter_checkpoint_path: str = ""
+    insightface_app_name: str = ""
+    insightface_providers: List = []
+
+    det_size_face: tuple = (640, 640)
+    target_size: int = 224
+    inference_steps: int = 20
+    guidance_scale: float = 7
+    ip_adapter_scale: float = 0.6
+    clip_skip: int = 2
+    num_images_per_prompt: int = 1

raySD/pydantic_model/__init__.py

@@ -0,0 +1,4 @@
+from .Image_Generate_Model import *
+from .Image_Response_Model import *
+from .Message_Model import *
+from .Stable_Diffusion_IP_Adapter_Model import *

raySD/utils/Enable_xFormers.py

@@ -0,0 +1,9 @@
+def enable_xformers_if_available(pipe):
+    try:
+        import xformers
+        pipe.enable_xformers_memory_efficient_attention()
+        print("xformers is available and has been enabled.")
+    except ImportError:
+        print("xformers is NOT installed. Running without xformers.")
+    except Exception as e:
+        print(f"xformers is installed but could not be enabled: {e}")

raySD/utils/__init__.py

@@ -0,0 +1 @@
+from .Enable_xFormers import *

raysd-0.1.0.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: raySD
+Version: 0.1.0
+Summary: Stable diffusion pipeline for ray
+Author: Nguyen Hong Ky
+Requires-Dist: numpy==1.26.4
+Requires-Dist: opencv-python>=4.10.0.84
+Requires-Dist: pillow>=9.5.0
+Requires-Dist: torch==2.3.1
+Requires-Dist: torchvision==0.18.1
+Requires-Dist: diffusers==0.33.0
+Requires-Dist: transformers==4.49.0
+Requires-Dist: peft==0.15.2
+Requires-Dist: openvino==2025.1.0
+Requires-Dist: ip-adapterv
+Requires-Dist: pydantic
+Requires-Dist: onnx==1.17.0
+Requires-Dist: onnxruntime==1.17.1
+Dynamic: author
+Dynamic: requires-dist
+Dynamic: summary

raysd-0.1.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+raySD/__init__.py,sha256=4ETmeXZJi5Myr8boHVhNqejNmMNHjG8oraKokGJw1ko,124
+raySD/pipeline/Abstract_Pipeline.py,sha256=0M0H7j0b7OrJpm-FD0ZCUqiv65DjxVn9YCd38TScbZg,768
+raySD/pipeline/Stable_Diffusion_IP_Adapter.py,sha256=0ViMt8X64TFIX9rZQT2ARFoExxN_FOfD5wUEJvoBQcQ,4564
+raySD/pipeline/Super_Resolution_Pipeline.py,sha256=5q2L5En7zMYdnKd-yUnkQyvQ4p3wyc1tu5laZaLAUkM,1470
+raySD/pipeline/__init__.py,sha256=3SohaavHvlq6mt5F5rpbReoU2pggZTH8T7BNE7GPPqo,84
+raySD/prompt/__init__.py,sha256=fVssgWdHZT0v98vMT328Cx_ayYSJh_RFmIhbts_0s8c,27
+raySD/prompt/chibi_prompt.py,sha256=u1v-1SSkdfM1QmSso99a1FcGdKW426eH3cRaGEbb_kE,4203
+raySD/pydantic_model/Image_Generate_Model.py,sha256=RlfMwF5GhWm6sQvxVN16HVv-nLth_B7yyrFvIvyX_WA,194
+raySD/pydantic_model/Image_Response_Model.py,sha256=tgZo2P-klG0TX_DRaReKHrH2ZUF-g53N1vnIMZLGNsw,98
+raySD/pydantic_model/Message_Model.py,sha256=D0DNP1F8kkwwB2X8QOTLtaLleVdf_INyIwjdypuVafM,94
+raySD/pydantic_model/Stable_Diffusion_IP_Adapter_Model.py,sha256=vpQM1To4K-amIKrFy0dxXxeLSMr5mWuAHJWWgdKGhO4,734
+raySD/pydantic_model/__init__.py,sha256=woiNl2YsEnyRvkFFRQuDIVLm-pMiJo27xgoRrHw0mOs,152
+raySD/utils/Enable_xFormers.py,sha256=2gZf3Shxth4jDkAtDFZa27wDFNavZa3BRn4nizgaiyw,392
+raySD/utils/__init__.py,sha256=KavMKUtooaiUDJOsNzdEUjOxTWxZ9sn8sUPWaJ2Mkgs,30
+raysd-0.1.0.dist-info/METADATA,sha256=YL59OSkrZJb-UxNft9_PTddMZGKH2LkmHiQeFBEs7Kg,598
+raysd-0.1.0.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+raysd-0.1.0.dist-info/top_level.txt,sha256=1yKh5zJFazWtmZ-WR7v930Nl7KkrjpodCSBFU76K7SU,6
+raysd-0.1.0.dist-info/RECORD,,

raysd-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.4.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

raysd-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+raySD