stablediffusion 3.5, deps update

Author: apocas

app/llms/ollamamultimodal.py

@@ -222,7 +222,7 @@ async def astream_chat(
 
 
 
-class OllamaMultiModal2(OllamaMultiModal):
+class OllamaMultiModalInternal(OllamaMultiModal):
     system: str = Field(
         default="", description="Default system message to send to the model."
     )

app/llms/workers/children/stablediffusion35.py

@@ -0,0 +1,48 @@
+import base64
+import io
+from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
+from diffusers import StableDiffusion3Pipeline
+import torch
+from transformers import T5EncoderModel
+
+from app.config import RESTAI_DEFAULT_DEVICE
+
+
+def worker(prompt, sharedmem):
+    model_id = "stabilityai/stable-diffusion-3.5-large-turbo"
+
+    nf4_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    model_nf4 = SD3Transformer2DModel.from_pretrained(
+        model_id,
+        subfolder="transformer",
+        quantization_config=nf4_config,
+        torch_dtype=torch.bfloat16
+    )
+
+    t5_nf4 = T5EncoderModel.from_pretrained("diffusers/t5-nf4", torch_dtype=torch.bfloat16)
+
+    pipeline = StableDiffusion3Pipeline.from_pretrained(
+        model_id, 
+        transformer=model_nf4,
+        text_encoder_3=t5_nf4,
+        torch_dtype=torch.bfloat16
+    )
+    pipeline.enable_model_cpu_offload()
+
+    image = pipeline(
+        prompt=prompt,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+        max_sequence_length=512,
+    ).images[0]
+
+
+    image_data = io.BytesIO()
+    image.save(image_data, format="JPEG")
+    image_base64 = base64.b64encode(image_data.getvalue()).decode('utf-8')
+
+    sharedmem["image"] = image_base64

app/llms/workers/stablediffusion.py

@@ -2,7 +2,7 @@
 
 #from app.llms.workers.children.stablediffusion import worker
 #from app.llms.workers.children.sdxl_lightning import worker
-from app.llms.workers.children.stablediffusion3 import worker
+from app.llms.workers.children.stablediffusion35 import worker
 
 try:
     set_start_method('spawn')

app/tools.py

@@ -29,9 +29,12 @@ def get_llm_class(llm_class_name):
     if llm_class_name == "Ollama":
         from app.llms.ollama import Ollama
         return Ollama, {}
-    elif llm_class_name == "OllamaMultiModal2":
-        from app.llms.ollamamultimodal import OllamaMultiModal2
-        return OllamaMultiModal2, {}
+    elif llm_class_name == "OllamaMultiModal":
+        from llama_index.multi_modal_llms.ollama import OllamaMultiModal
+        return OllamaMultiModal, {}
+    elif llm_class_name == "OllamaMultiModalInternal" or llm_class_name == "OllamaMultiModal2":
+        from app.llms.ollamamultimodal import OllamaMultiModalInternal
+        return OllamaMultiModalInternal, {}
     elif llm_class_name == "OpenAI":
         from llama_index.llms.openai import OpenAI
         return OpenAI, {}