2609

src/sample.ipynb

@@ -30,15 +30,17 @@
     "    \"AiArtLab/sdxs3d\", subfolder=\"vae\", torch_dtype=dtype\n",
     ").to(device).eval()\n",
     "\n",
-    "unet = UNet2DConditionModel.from_pretrained( \"/workspace/sdxs3d/unet\"#\"AiArtLab/sdxs3d\"\n",
-    "                                            , subfolder=\"unet\", torch_dtype=dtype\n",
+    "unet = UNet2DConditionModel.from_pretrained(\n",
+    "    \"AiArtLab/sdxs3d\" \n",
+    "    #\"/workspace/sdxs3d\" \n",
+    "    , subfolder=\"unet\", torch_dtype=dtype\n",
     ").to(device).eval()\n",
     "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-Embedding-0.6B\", padding_side=\"left\")\n",
-    "text_model = AutoModel.from_pretrained(\"Qwen/Qwen3-Embedding-0.6B\").to(device).eval()\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\"AiArtLab/sdxs3d\", subfolder=\"tokenizer\")\n",
+    "text_model = AutoModel.from_pretrained(\"AiArtLab/sdxs3d\", subfolder=\"text_encoder\").to(device).eval()\n",
     "\n",
     "# ====== FlowMatch Scheduler ======\n",
-    "scheduler = FlowMatchEulerDiscreteScheduler()\n",
+    "scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(\"AiArtLab/sdxs3d\", subfolder=\"scheduler\")\n",
     "print('loaded')\n",
     "\n"
    ]
@@ -271,7 +273,7 @@
     "    generator = torch.Generator(device=device).manual_seed(42)\n",
     ")\n",
     "\n",
-    "grid = display_image_grid(images,prompts, cols=3, save_path=\"../result_grid.jpg\")\n"
+    "grid = display_image_grid(images,prompts, cols=3, save_path=\"result_grid.jpg\")\n"
    ]
   },
   {
@@ -313,7 +315,266 @@
    "id": "b08fbf66-8bd1-4a20-8715-0e748a07a932",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "import gradio as gr\n",
+    "import numpy as np\n",
+    "import random\n",
+    "\n",
+    "import spaces #[uncomment to use ZeroGPU]\n",
+    "import torch\n",
+    "\n",
+    "from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel, FlowMatchEulerDiscreteScheduler\n",
+    "from transformers import AutoTokenizer, AutoModel\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "model_repo_id = \"AiArtLab/sdxs3d\"  # Replace to the model you would like to use\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    dtype = torch.float16\n",
+    "else:\n",
+    "    dtype = torch.float32\n",
+    "\n",
+    "\n",
+    "class SimpleDiffusionPipeline(DiffusionPipeline):\n",
+    "    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler):\n",
+    "        super().__init__()\n",
+    "        self.register_modules(\n",
+    "            vae=vae,\n",
+    "            text_encoder=text_encoder,\n",
+    "            tokenizer=tokenizer,\n",
+    "            unet=unet,\n",
+    "            scheduler=scheduler,\n",
+    "        )\n",
+    "\n",
+    "    @torch.no_grad()\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        prompt,\n",
+    "        negative_prompt=None,\n",
+    "        height=512,\n",
+    "        width=512,\n",
+    "        num_inference_steps=50,\n",
+    "        guidance_scale=4.0,\n",
+    "        generator=None,\n",
+    "        **kwargs,\n",
+    "    ):\n",
+    "        batch_size = len(prompt) if isinstance(prompt, list) else 1\n",
+    "\n",
+    "        # 1. Токенизация\n",
+    "        toks = self.tokenizer(\n",
+    "            prompt,\n",
+    "            padding=\"max_length\",\n",
+    "            truncation=True,\n",
+    "            max_length=512,\n",
+    "            return_tensors=\"pt\"\n",
+    "        ).to(self.device)\n",
+    "\n",
+    "        outs = self.text_encoder(**toks)\n",
+    "        text_emb = outs.last_hidden_state[:, -1].unsqueeze(1)  # твой last_token_pool\n",
+    "\n",
+    "        if negative_prompt is not None:\n",
+    "            neg_toks = self.tokenizer(\n",
+    "                negative_prompt,\n",
+    "                padding=\"max_length\",\n",
+    "                truncation=True,\n",
+    "                max_length=512,\n",
+    "                return_tensors=\"pt\"\n",
+    "            ).to(self.device)\n",
+    "            neg_outs = self.text_encoder(**neg_toks)\n",
+    "            neg_emb = neg_outs.last_hidden_state[:, -1].unsqueeze(1)\n",
+    "        else:\n",
+    "            neg_emb = torch.zeros_like(text_emb)\n",
+    "\n",
+    "        # guidance\n",
+    "        if guidance_scale != 1.0:\n",
+    "            text_emb = torch.cat([neg_emb, text_emb])\n",
+    "\n",
+    "        # 2. Латенты\n",
+    "        latents = torch.randn(\n",
+    "            (batch_size, self.unet.config.in_channels, height // self.vae.config.scaling_factor, width // self.vae.config.scaling_factor),\n",
+    "            device=self.device,\n",
+    "            dtype=torch.float16,\n",
+    "            generator=generator,\n",
+    "        )\n",
+    "\n",
+    "        self.scheduler.set_timesteps(num_inference_steps, device=self.device)\n",
+    "\n",
+    "        # 3. Диффузия\n",
+    "        for t in self.scheduler.timesteps:\n",
+    "            latent_input = torch.cat([latents, latents]) if guidance_scale != 1.0 else latents\n",
+    "            flow = self.unet(latent_input, t, encoder_hidden_states=text_emb).sample\n",
+    "\n",
+    "            if guidance_scale != 1.0:\n",
+    "                flow_uncond, flow_cond = flow.chunk(2)\n",
+    "                flow = flow_uncond + guidance_scale * (flow_cond - flow_uncond)\n",
+    "\n",
+    "            latents = self.scheduler.step(flow, t, latents).prev_sample\n",
+    "\n",
+    "        # 4. Декод\n",
+    "        latents = latents / self.vae.config.scaling_factor\n",
+    "        images = self.vae.decode(latents).sample\n",
+    "        images = (images / 2 + 0.5).clamp(0, 1)\n",
+    "\n",
+    "        return images\n",
+    "\n",
+    "\n",
+    "vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder=\"vae\", torch_dtype=dtype).to(device)\n",
+    "unet = UNet2DConditionModel.from_pretrained(model_repo_id, subfolder=\"unet\", torch_dtype=dtype).to(device)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_repo_id, subfolder=\"tokenizer\")\n",
+    "text_encoder = AutoModel.from_pretrained(model_repo_id, subfolder=\"text_encoder\", torch_dtype=dtype).to(device)\n",
+    "scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_repo_id, subfolder=\"scheduler\")\n",
+    "\n",
+    "pipe = SimpleDiffusionPipeline(\n",
+    "    vae=vae,\n",
+    "    text_encoder=text_encoder,\n",
+    "    tokenizer=tokenizer,\n",
+    "    unet=unet,\n",
+    "    scheduler=scheduler,\n",
+    ").to(device)\n",
+    "\n",
+    "\n",
+    "MAX_SEED = np.iinfo(np.int32).max\n",
+    "MAX_IMAGE_SIZE = 384\n",
+    "\n",
+    "\n",
+    "@spaces.GPU #[uncomment to use ZeroGPU]\n",
+    "def infer(\n",
+    "    prompt,\n",
+    "    negative_prompt,\n",
+    "    seed,\n",
+    "    randomize_seed,\n",
+    "    width,\n",
+    "    height,\n",
+    "    guidance_scale,\n",
+    "    num_inference_steps,\n",
+    "    progress=gr.Progress(track_tqdm=True),\n",
+    "):\n",
+    "    if randomize_seed:\n",
+    "        seed = random.randint(0, MAX_SEED)\n",
+    "\n",
+    "    generator = torch.Generator(device=device).manual_seed(seed)  # ← используйте seed, а не 42!\n",
+    "\n",
+    "    # Генерация\n",
+    "    images_tensor = pipe(\n",
+    "        prompt=prompt,\n",
+    "        negative_prompt=negative_prompt,\n",
+    "        guidance_scale=guidance_scale,\n",
+    "        num_inference_steps=num_inference_steps,\n",
+    "        width=width,\n",
+    "        height=height,\n",
+    "        generator=generator,\n",
+    "    )  # [B, C, H, W]\n",
+    "\n",
+    "    # Конвертация в numpy для Gradio\n",
+    "    image = images_tensor[0].cpu().permute(1, 2, 0).numpy()\n",
+    "    image = (image * 255).astype(np.uint8)\n",
+    "\n",
+    "    return image, seed\n",
+    "\n",
+    "\n",
+    "examples = [\n",
+    "    \"A delicious ceviche cheesecake slice\",\n",
+    "    \"ариец в имперских доспехах будущего\",\n",
+    "    \"A close-up image of an astronaut's helmet with a frosted and opaque visor. The visor reflects the cold, frozen texture of space. Resting on the surface of the visor is a butterfly with vibrant, intricately patterned wings. The contrast between the delicate natural beauty of the butterfly and the cold, industrial helmet creates a striking image. The butterfly adds a touch of fragility and life to the otherwise harsh and unfeeling setting. The faint glow of distant stars can be seen through the frost, further enhancing the surreal atmosphere.\",    \n",
+    "]\n",
+    "\n",
+    "css = \"\"\"\n",
+    "#col-container {\n",
+    "    margin: 0 auto;\n",
+    "    max-width: 640px;\n",
+    "}\n",
+    "\"\"\"\n",
+    "\n",
+    "with gr.Blocks(css=css) as demo:\n",
+    "    with gr.Column(elem_id=\"col-container\"):\n",
+    "        gr.Markdown(\" # Text-to-Image Gradio Template\")\n",
+    "\n",
+    "        with gr.Row():\n",
+    "            prompt = gr.Text(\n",
+    "                label=\"Prompt\",\n",
+    "                show_label=False,\n",
+    "                max_lines=1,\n",
+    "                placeholder=\"Enter your prompt\",\n",
+    "                container=False,\n",
+    "            )\n",
+    "\n",
+    "            run_button = gr.Button(\"Run\", scale=0, variant=\"primary\")\n",
+    "\n",
+    "        result = gr.Image(label=\"Result\", show_label=False)\n",
+    "\n",
+    "        with gr.Accordion(\"Advanced Settings\", open=False):\n",
+    "            negative_prompt = gr.Text(\n",
+    "                label=\"Negative prompt\",\n",
+    "                max_lines=1,\n",
+    "                placeholder=\"Enter a negative prompt\",\n",
+    "                visible=True,\n",
+    "                value =\"low quality\"\n",
+    "            )\n",
+    "\n",
+    "            seed = gr.Slider(\n",
+    "                label=\"Seed\",\n",
+    "                minimum=0,\n",
+    "                maximum=MAX_SEED,\n",
+    "                step=1,\n",
+    "                value=0,\n",
+    "            )\n",
+    "\n",
+    "            randomize_seed = gr.Checkbox(label=\"Randomize seed\", value=True)\n",
+    "\n",
+    "            with gr.Row():\n",
+    "                width = gr.Slider(\n",
+    "                    label=\"Width\",\n",
+    "                    minimum=192,\n",
+    "                    maximum=MAX_IMAGE_SIZE,\n",
+    "                    step=64,\n",
+    "                    value=256,  # Replace with defaults that work for your model\n",
+    "                )\n",
+    "\n",
+    "                height = gr.Slider(\n",
+    "                    label=\"Height\",\n",
+    "                    minimum=192,\n",
+    "                    maximum=MAX_IMAGE_SIZE,\n",
+    "                    step=64,\n",
+    "                    value=384,  # Replace with defaults that work for your model\n",
+    "                )\n",
+    "\n",
+    "            with gr.Row():\n",
+    "                guidance_scale = gr.Slider(\n",
+    "                    label=\"Guidance scale\",\n",
+    "                    minimum=0.0,\n",
+    "                    maximum=10.0,\n",
+    "                    step=0.1,\n",
+    "                    value=4.0,  # Replace with defaults that work for your model\n",
+    "                )\n",
+    "\n",
+    "                num_inference_steps = gr.Slider(\n",
+    "                    label=\"Number of inference steps\",\n",
+    "                    minimum=1,\n",
+    "                    maximum=50,\n",
+    "                    step=1,\n",
+    "                    value=40,  # Replace with defaults that work for your model\n",
+    "                )\n",
+    "\n",
+    "        gr.Examples(examples=examples, inputs=[prompt])\n",
+    "    gr.on(\n",
+    "        triggers=[run_button.click, prompt.submit],\n",
+    "        fn=infer,\n",
+    "        inputs=[\n",
+    "            prompt,\n",
+    "            negative_prompt,\n",
+    "            seed,\n",
+    "            randomize_seed,\n",
+    "            width,\n",
+    "            height,\n",
+    "            guidance_scale,\n",
+    "            num_inference_steps,\n",
+    "        ],\n",
+    "        outputs=[result, seed],\n",
+    "    )\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    demo.launch()"
+   ]
   }
  ],
  "metadata": {

unet/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:afc06beff07034f0ce9f671c83222e7f78eedc3b3ce93293143accdebef1b111
-size 1887
+oid sha256:0ef8fbaff98c8d479d68b566d07ef4fb8e51ac26b9e8b5a3cb2b23f9a978f6ca
+size 1874

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88a4685341dba55274c8ae3991144a18be96dd6257c81612e749a408e934544f
-size 3092571208
+oid sha256:e069f7e9f439bba567cd93aa9942ed3481c57a542dceed41fa78f9c97a344dfe
+size 6184944280